[peterdsharpe]

Problem

In AsyncBatchedCollector's thread worker (_env_loop), pool.async_step_and_maybe_reset_recv returns a TensorDict where "env_index" is a NonTensorData(env_index) marker set by the pool, and the worker overwrites it with a Python int:

cur_td.set(_ENV_IDX_KEY, env_id)

In tensordict 0.15-0.16, writing a Python int over an existing NonTensorData entry converts it to a 0-D tensor and emits a FutureWarning. Starting in tensordict 0.17, set preserves the existing entry kind and keeps "env_index" as NonTensorData, which breaks downstream consumers that expect an integer tensor (such as ReplayBufferEnsemble routing and .unique()).

Fix

Write an explicit torch.long tensor matching cur_td.batch_size and cur_td.device, the same way the chunked process-worker path in _process_env_loop already does:

cur_td.set(
    _ENV_IDX_KEY,
    torch.full(cur_td.batch_size, env_id, dtype=torch.long, device=cur_td.device),
)

For an unbatched env slot this produces the same 0-D int64 tensor as today without triggering the deprecation warning.

Tests

Added TestAsyncBatchedCollector::test_thread_worker_env_index_is_long_tensor in test/test_inference_server.py, which runs a threaded AsyncBatchedCollector with the FutureWarning turned into an error and checks that "env_index" is a torch.long tensor.

Related: #4551, #4576, and #4638.

[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.

❌ 1 New Failure

As of commit 0ca7025 with merge base 6479319 ():

NEW FAILURE - The following job has failed:

• Unit-tests on Linux / tests-cpu (3.10, bulk) / linux-job (gh)
  test/modules/test_rnn.py::TestLSTMModule::test_single_step[False-shape3]

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

[bsprenger]

Looks great, thanks for fixing, and sorry for the delay in reviewing! Can you rebase + rerun CI before we merge?