[peterdsharpe]

With tensordict 0.17, AsyncBatchedCollector thread workers yield batches whose env_index is non-tensor data instead of a long tensor. Anything that routes or groups transitions by that index then breaks: in a run simulating 0.17's behavior, 5 TestAsyncBatchedCollector tests fail, because ReplayBufferEnsemble needs integer member ids and .unique() does not work on a non-tensor entry. TorchRL pins tensordict<0.16, so no released TorchRL is affected, but this has to land before the pin is lifted.

The fix stores env_index in the thread worker loop as a long tensor of the result's batch size, on the result's device, the same way the chunked process workers already do.

Details

In _env_loop, each result comes from AsyncEnvPool.async_step_and_maybe_reset_recv, which for the queue exchange marks it with NonTensorData(data=env_index). The loop then wrote the int env_id over that entry. tensordict 0.15 and 0.16 cast an int written over a NonTensorData entry to a 0-d tensor and emit a FutureWarning; from 0.17 the int stays NonTensorData, like the entry it replaces.

Now the loop writes torch.full(cur_td.batch_size, env_id, dtype=torch.long, device=cur_td.device). For an unbatched slot this is the same 0-d int64 tensor the cast gives today, so on current tensordict the only change is that the warning goes away. A pool slot with a non-empty batch size now gets one index per row; today the cast raises in that case.

The other two places that set env_index are fine. The chunk_size == 1 process loop writes the int on a key that is not set yet, which already gives a tensor and is not affected by the 0.17 change. The chunked path already uses torch.full(..., dtype=torch.long).

Tests

I added TestAsyncBatchedCollector::test_thread_worker_env_index_is_long_tensor. It turns the tensordict FutureWarning into an error and checks that a thread-worker batch's env_index is a long tensor with values in {0, 1}. Without the filter the test would pass on current tensordict, since tensordict casts the int; the filter is what catches the old write.

Run on a 24-core desktop (CPU, torch 2.14.1, Python 3.12, tensordict main):

• pytest test/test_inference_server.py -k test_thread_worker_env_index_is_long_tensor: fails without the fix (the worker thread raises the FutureWarning, and the collector reports RuntimeError: A collector worker thread raised an exception.), passes with it.
• pytest test/test_inference_server.py -k TestAsyncBatchedCollector -n 4: 106 passed, 6 skipped.

I also ran this branch together with #4638 against a copy of tensordict main patched to do what 0.17 announces (the removals raise, the new defaults apply): test_collectors.py, test_inference_server.py, test_dt.py, test_exploration.py and test/llm give the same results as on unpatched tensordict main, and none of the removed APIs is reached.

Related: #4551 (tensordict 0.16) and #4576 (tensordict 0.17).

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/4637

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 2 New Failures

As of commit 39144dd with merge base 380dd55 ():

NEW FAILURES - The following jobs have failed:

• Unit-tests on Linux / tests-cpu (3.10, bulk) / linux-job (gh)
  test/test_configs.py::TestHydraParsing::test_transformed_env_parsing_with_file
• Unit-tests on Linux / tests-cpu (3.13, bulk) / linux-job (gh)
  test/test_configs.py::TestHydraParsing::test_transformed_env_parsing_with_file

This comment was automatically generated by Dr. CI and updates every 15 minutes.