Add TorchRL environment wrapper for isaaclab_rl

[theap06]

Description

Adds isaaclab_rl.torchrl.IsaacLabTorchRLWrapper, a torchrl.envs.EnvBase implementation over Isaac Lab's batched environments, together with a torchrl extra, API docs, and tests.

Validated on an A100-40GB (Isaac Sim 6.0, torchrl 0.13.3, 64 envs per task):

Check	Result
test_torchrl_wrapper_specs.py (kit-free)	passed
test_torchrl_wrapper.py (first 5 registered tasks)	2 passed
Same checks over Ant, Ant-Direct, Cartpole, Cartpole-Direct, Cartpole-Camera, Cartpole-Camera-Direct, Humanoid-Direct, Velocity-Flat-AnymalD, Reorient-Cube-Shadow-OpenAI-FF-Direct	9/9 OK, no NaNs, terminal observation reported on every done row
Cartpole, 200 steps, vs torchrl.envs.libs.isaac_lab.IsaacLabWrapper	12.5k env-steps/s with terminal observations on done rows; upstream 9.6k (native_autoreset=True) / 10.4k (False) with NaN next-observations on done rows

[greptile-apps]

Greptile Summary

This PR adds an optional TorchRL dependency and a TorchRL EnvBase adapter for batched Isaac Lab environments, including structured specifications, terminal-observation handling, documentation, and tests.

• Converts Gymnasium observation and action spaces into batched TorchRL specifications.
• Preserves Isaac Lab same-step autoreset behavior and terminal observations in TorchRL rollouts.
• Adds kit-free contract tests and simulator-backed rollout coverage.
• The constructor currently admits non-RL manager-based environments that cannot satisfy the wrapper contract.

Confidence Score: 4/5

The PR should not merge until the constructor stops accepting non-RL ManagerBasedEnv variants or fully supports their different interface.

The new type gate promises support for two environment classes that lack the spaces and five-value step result the wrapper unconditionally requires, causing construction or stepping to fail for those accepted inputs.

Files Needing Attention: source/isaaclab_rl/isaaclab_rl/torchrl/vecenv_wrapper.py

Important Files Changed

Filename	Overview
source/isaaclab_rl/isaaclab_rl/torchrl/vecenv_wrapper.py	Adds the TorchRL adapter and specification conversion, but its accepted-type check includes non-RL manager-based environments that fail the required construction and step contracts.
source/isaaclab_rl/test/test_torchrl_wrapper_specs.py	Thoroughly tests specifications, buffer cloning, terminal observations, reset behavior, clipping, device placement, and lifecycle using an RL-contract fake, but does not cover the newly accepted non-RL manager-based classes.
source/isaaclab_rl/test/test_torchrl_wrapper.py	Adds simulator-backed rollouts validating TorchRL specifications, finite-horizon termination behavior, and terminal observations across registered RL tasks.
pyproject.toml	Adds TorchRL 0.13 or newer as an isolated optional dependency.
docs/source/api/lab_rl/isaaclab_rl.rst	Publishes the new TorchRL wrapper in the isaaclab_rl API documentation.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  T[TorchRL TensorDict action] --> W[IsaacLabTorchRLWrapper]
  W --> E[Isaac Lab batched RL environment]
  E --> S[Observation, reward, termination, truncation, extras]
  S --> F{Final observation available?}
  F -->|Yes| O[Use terminal observation on done rows]
  F -->|No| N[Invalidate done-row observations]
  O --> R[TorchRL next TensorDict]
  N --> R
  R --> A[Masked autoreset request]
  A --> C[Return current post-reset Isaac Lab observations]

Loading

_{Reviews (1): Last reviewed commit: "Add Kit-dependent integration test for t..." | Re-trigger Greptile}

[isaaclab-review-bot]

Isaac Lab Review Bot

Adds a direct TorchRL EnvBase wrapper with specs, auto-reset handling, optional dependency metadata, documentation, and tests. Two contract mismatches need correction: the constructor accepts non-RL manager environments that _step cannot consume, and actions are not transferred back to the simulation device when the wrapper uses a different TensorDict device.

• Design and architecture: The direct EnvBase integration preserves batched observation groups and explicitly handles Isaac Lab’s same-step auto-reset behavior. However, the accepted-type set includes ManagerBasedEnv and its Warp counterpart even though the implementation requires the five-value RL step() contract; construction should be restricted to supported RL environment classes.
• API: The public wrapper is exported, documented, and backed by a dedicated torchrl extra. The constructor’s runtime acceptance currently conflicts with its annotation and error message by permitting non-RL manager environments, while the documented device option implies cross-device TensorDict support that the action path does not fully implement.
• Implementation: Spec conversion, buffer cloning, finite-horizon handling, and terminal-observation behavior are covered by focused tests. Before stepping, actions must be moved from the wrapper device to self.unwrapped.device; otherwise a non-default wrapper device produces device-mismatched simulator inputs.

Minor fixes needed. Posted 2 actionable findings inline.

Automated review; human maintainers own approval decisions.

[isaaclab-review-bot]

🟡 Warning · Implementation — Actions not moved to simulation device

The docstring advertises device as the device the tensordicts live on, and outputs are converted via the TensorDict(device=...) construction. Actions, however, are forwarded to self.env.step unchanged, so a wrapper constructed with a device different from self.unwrapped.device (e.g. device="cpu" over a CUDA env) hands device-mismatched tensors to the simulator. Convert actions to the environment device before stepping.

[theap06]

The supported env classes already move the action to their own device at the top of step(): DirectRLEnv.step does action = action.to(self.device), ManagerBasedRLEnv.step calls process_action(action.to(self.device)), and DirectRLEnvWarp / ManagerBasedRLEnvWarp do the same before the Warp copy. A wrapper constructed with device="cpu" over a CUDA env therefore steps correctly as is, so I left this unchanged.

[theap06]

I left it here just for performance reasons; however, open to migrating this to be global

[kellyguo11]

thanks for the PR! do you happen to have an example cfg with torch RL that can be used for testing? maybe a cartpole training configuration?

[theap06]

@kellyguo11 yeah! I'm running some cartpole ppo experiments. I also wanted to ask if it makes sense to keep this as a standalone env or subclass of torchRL's IsaacLabWrapper.

[theap06]

@kellyguo11 this is the PR for the IsaacLab Wrapper pytorch/rl#2937

[kellyguo11]

I think a standalone env class would make sense so that we can align it with the Isaac Lab design more closely and will be easier for us to maintain. would be good to have the cartpole example along with the PR so that we can review the full training setup.

[theap06]

@kellyguo11 ran a ppo experiment with cartpole and torchrl learns faster and its final policy balances the pole more reliably. The reward curve is more noisy later in training though.

[theap06]

The only difference in the experiment setups which is what likely caused the noise was that RSL-RL uses adaptive kl while torchRL here just uses ppo loss. I can change that parameter setup, but the curves demonstrate that the wrapper works soundly with isaac lab

[kellyguo11]

Thanks, we'll let the team do a more thorough review of the changes.

[AntoineRichard]

AI-generated review (Codex)

This review was produced by an AI coding agent and should be validated by a human maintainer.

I would request changes before merging.

Findings

• Important — pyproject.toml:156, test_torchrl_wrapper_specs.py:20: torchrl was added to pyproject.toml but not uv.lock or the CI environment. uv export --frozen --extra torchrl reports that the extra is undefined in the locked metadata, and both new test suites use importorskip, so standard CI skips all TorchRL coverage. Regenerate uv.lock and install the extra in the isaaclab_rl test lane.

• Important — vecenv_wrapper.py:196, train_torchrl.py:106: converted DirectMARLEnv tasks never receive terminal observations. The MARL environment stores them under extras[agent]["final_obs"], while multi_agent_to_single_agent passes that structure unchanged and the wrapper only checks top-level extras["final_obs"]. Timeouts therefore produce NaN next-observations instead of values suitable for PPO bootstrapping. Convert per-agent terminal observations alongside normal observations and add a MARL-timeout regression test.

• Important — vecenv_wrapper.py:141: the public device option converts outputs but not inputs. With a CUDA simulation wrapped using device="cpu", TorchRL produces CPU actions and _step forwards them directly into CUDA task code. Move actions to self.unwrapped.device before stepping and test cross-device operation.

• Important — ppo.py:29: the exported PPO API accepts any IsaacLabTorchRLWrapper, but assumes flat continuous actions and flat "policy" observations. A (2, 3) Box produces (N, 3) actions rather than (N, 2, 3); Discrete/MultiDiscrete spaces receive Gaussian floats; Composite policy observations reach an MLP expecting a tensor. Either implement flattening/reshaping and categorical policies or reject unsupported specs explicitly.

• Important — ppo.py:45: when a "critic" group is Composite, make_critic silently switches to "policy", discarding privileged observations. This contradicts ppo_cfg.py:33. Flatten/encode Composite critic observations or raise an unsupported-space error.

• Important — vecenv_wrapper.py:196: extras["final_obs"] persists between steps, but the wrapper treats key presence as validity for every current done row. Visualizer-triggered resets occur after normal terminal-observation capture, so they can receive a stale terminal observation from an earlier episode. Clear/carry a validity mask per step and capture manual-reset observations before reset.

• Important — tools/environ_docs.py:52: the generated environment browser was not refreshed. tools/update_environments_rst.py --check fails, and regeneration adds TorchRL to both Cartpole rows. The browser also lacks a torchrl extra mapping and enables playback even though dispatch.py:30 has no TorchRL play backend. Commit the generated rows, add the extra mapping, and either implement playback or make the UI/docs explicitly train-only.

• Important — train_torchrl.py:116, ppo.py:103: KeyboardInterrupt unwinds train_ppo before the collector and TensorBoard writer are closed, then is suppressed and reported like successful completion. Other exceptions additionally bypass environment cleanup. Put collector, writer, and environment cleanup in finally blocks and report interruptions distinctly.

Validation

• git diff --check: passed.
• Changelog validation: passed.
• Root pyproject tests: 11 passed.
• Environment documentation check: failed as described.
• TorchRL unit suite: skipped because TorchRL is absent.
• Full simulator tests were not rerun.