[Feature] Add functorch integration tests for TensorDictModule - Fixes #154

[ParamThakkar123]

Summary

Adds a comprehensive test suite for functorch integration with TensorDictModule to address issue #154.

Tests Added

• test_tdmodule_functional_params - Tests extracting functional params from a TDModule using TensorDict.from_module
• test_vmap_on_tdmodule_with_params - Tests vmap on TensorDictModule with functional params expanded across batch dimension
• test_nested_tdmodule_param_length - Tests nested TDModules (SafeProbabilisticTensorDictSequential) param extraction
• test_nested_tdmodule_param_casting - Tests that params can be cast to different dtypes (e.g., float64)
• test_tdsequence_vmap_params - Tests TDSequence param handling with vmap
• test_vmap_multiple_inputs - Tests vmap with multiple input keys to module
• test_vmap_module_class - Tests VmapModule class directly
• test_tdsequence_vmap - Tests vmap on SafeSequential
• test_tdmodule_functional_to_module - Tests the to_module context manager for functional params
• test_nested_sequential_vmap - Tests vmap on nested TensorDictSequential modules

All tests cover the functionality requested in the issue.

Fixes #154

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

⚠️ 17 Awaiting Approval

As of commit 1d50328 with merge base c764978 ():

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 (gh)
• Continuous Benchmark (PR) (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)
• Unit-tests on Windows (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: Add functorch integration tests for TensorDictModule - Fixes #154

Supported Prefixes (case-sensitive)

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

Prefix	Label Applied	Example
[BugFix]	BugFix	[BugFix] Fix memory leak in collector
[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
[Environment] or [Environments]	Environments	[Environments] Add Gymnasium support
[Data]	Data	[Data] Fix replay buffer sampling
[Performance] or [Perf]	Performance	[Performance] Optimize tensor ops
[BC-Breaking]	bc breaking	[BC-Breaking] Remove deprecated API
[Deprecation]	Deprecation	[Deprecation] Mark old function
[Quality]	Quality	[Quality] Fix typos and add codespell

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