[KhushiKeswani]

Description

• Fix IncrementalTokenizer._step() raising a RuntimeError when accessing tokens.full from a LazyStackedTensorDict containing sequences of different lengths.
• Use get(..., as_list=True) for LazyStackedTensorDict to retrieve individual token sequences without attempting to stack tensors with incompatible shapes.
• Preserve the existing retrieval behavior for regular TensorDict instances.
• Add regression tests covering both cases:
  • Unequal token lengths when history is absent.
  • Unequal token lengths when existing history is extended.

Motivation and Context

When processing a 1-D batch with variable-length token sequences, LazyStackedTensorDict.get() attempts to stack the individual tensors. Since their lengths differ, this raises a RuntimeError.

This change retrieves the sequences individually for lazy stacks while preserving the existing behavior for regular TensorDicts.

Fixes #4686.

• I have raised an issue to propose this change ([required](https://github.com/pytorch/rl/issues) for new features and bug fixes)

Types of changes

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds core functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Documentation (update in the documentation)
• Example (update in the folder of examples)

Checklist

• I have read the [CONTRIBUTION](https://github.com/pytorch/rl/blob/main/CONTRIBUTING.md) guide (required)
• My change requires a change to the documentation.
• I have updated the tests accordingly (required for a bug fix or a new feature).
• I have updated the documentation accordingly.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

⚠️ 16 Awaiting Approval

As of commit ebc8606 with merge base a711ef9 ():

AWAITING APPROVAL - The following workflows need approval before CI can run:

• Build Aarch64 Linux Wheels (gh)
• Build Linux Wheels (gh)
• Build M1 Wheels (gh)
• Build Windows Wheels (gh)
• Continuous Benchmark (PR) (gh)
• Examples Tests on Linux (gh)
• Generate documentation (gh)
• Habitat Tests on Linux (gh)
• Libs Tests on Linux (gh)
• Lint (gh)
• LLM Tests on Linux (gh)
• Push Binary Nightly (gh)
• SOTA Tests on Linux (gh)
• Tutorials Tests on Linux (gh)
• Unit-tests on Linux (gh)
• Validate Test Partitioning (gh)

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

[github-actions]

⚠️ PR Title Label Error

PR title must start with a label prefix in brackets (e.g., [BugFix]).

Current title: Fix IncrementalTokenizer with unequal token lengths

Supported Prefixes (case-sensitive)

Your PR title must start with exactly one of these prefixes:

Prefix	Label Applied	Example
[Algorithm]	new algo	[Algorithm] Add new RL objective
[BE]	BE	[BE] Improve error messages
[Benchmark] or [Benchmarks]	Benchmarks	[Benchmark] Add collector benchmark
[BugFix]	BugFix	[BugFix] Fix memory leak in collector
[Example] or [Examples]	Examples	[Example] Add training script
[Feature]	Feature	[Feature] Add new optimizer
[Doc] or [Docs]	Documentation	[Doc] Update installation guide
[Refactor]	Refactoring	[Refactor] Clean up module imports
[CI]	CI	[CI] Fix workflow permissions
[Test] or [Tests]	Tests	[Tests] Add unit tests for buffer
[Trainer] or [Trainers]	Trainers	[Trainer] Add trainer config
[Environment] or [Environments]	Environments	[Environments] Add Gymnasium support
[Data]	Data	[Data] Fix replay buffer sampling
[LLM]	llm/	[LLM] Add reward model integration
[Minor]	small change	[Minor] Fix typo in error message
[Performance] or [Perf]	Performance	[Performance] Optimize tensor ops
[BC-Breaking]	bc breaking	[BC-Breaking] Remove deprecated API
[Deprecation]	Deprecation	[Deprecation] Mark old function
[Algorithm] or [Algorithms]	new algo	[Algorithm] Add new objective
[Quality]	Quality	[Quality] Fix typos and add codespell
[Versioning]	versioning	[Versioning] Bump release version
[WIP]	WIP	[WIP] Draft implementation

Note: Common variations like singular/plural are supported (e.g., [Doc] or [Docs]).

[github-actions]

⚠️ PR Title Label Error

Unknown or invalid prefix [BUGFIX].

Current title: [BUGFIX] Fix IncrementalTokenizer with unequal token lengths

Supported Prefixes (case-sensitive)

Your PR title must start with exactly one of these prefixes:

Prefix	Label Applied	Example
[Algorithm]	new algo	[Algorithm] Add new RL objective
[BE]	BE	[BE] Improve error messages
[Benchmark] or [Benchmarks]	Benchmarks	[Benchmark] Add collector benchmark
[BugFix]	BugFix	[BugFix] Fix memory leak in collector
[Example] or [Examples]	Examples	[Example] Add training script
[Feature]	Feature	[Feature] Add new optimizer
[Doc] or [Docs]	Documentation	[Doc] Update installation guide
[Refactor]	Refactoring	[Refactor] Clean up module imports
[CI]	CI	[CI] Fix workflow permissions
[Test] or [Tests]	Tests	[Tests] Add unit tests for buffer
[Trainer] or [Trainers]	Trainers	[Trainer] Add trainer config
[Environment] or [Environments]	Environments	[Environments] Add Gymnasium support
[Data]	Data	[Data] Fix replay buffer sampling
[LLM]	llm/	[LLM] Add reward model integration
[Minor]	small change	[Minor] Fix typo in error message
[Performance] or [Perf]	Performance	[Performance] Optimize tensor ops
[BC-Breaking]	bc breaking	[BC-Breaking] Remove deprecated API
[Deprecation]	Deprecation	[Deprecation] Mark old function
[Algorithm] or [Algorithms]	new algo	[Algorithm] Add new objective
[Quality]	Quality	[Quality] Fix typos and add codespell
[Versioning]	versioning	[Versioning] Bump release version
[WIP]	WIP	[WIP] Draft implementation

Note: Common variations like singular/plural are supported (e.g., [Doc] or [Docs]).

[coder-jayp]

Thanks for the contribution

Left inline comments—the rollout from #4686 still raises RuntimeError during env.step() because only the child "tokens" is lazy-stacked while the root is a standard TensorDict.

[coder-jayp]

In ChatEnv, the root action_td is a normal TensorDict and only the child "tokens" is lazy-stacked. The isinstance check evaluates to False, falls back to else:, and env.step(action_td) still crashes with RuntimeError.

Using as_list=True directly in tensordict.get() handles this without needing the isinstance check.

[coder-jayp]

With as_list=True, existing_tokens_full becomes a Python list. Setting it directly stores it as NonTensorData (which is why .data was needed in the test assertion). It should be converted back into a nested tensor with torch.nested.as_nested_tensor before being set on next_tensordict.

[coder-jayp]

Could this test be updated to call the public env.step(action_td) directly on the output of env.reset()? Calling private _step() with a manual lazy_stack bypasses how ChatEnv actually packages tensordicts in real rollouts