[Feature] Surface Isaac Lab per-index reset and reset_to via IsaacLabWrapper

docs/source/reference/envs_api.rst

@@ -191,6 +191,19 @@ With these, the following methods are implemented:
   a :class:`tensordict.TensorDict` input. It return the first tensordict of a rollout, usually
   containing a ``"done"`` state and a set of observations. If not present,
   a ``"reward"`` key will be instantiated with 0s and the appropriate shape.
+
+  To reset *deterministically* to a state contained in the input tensordict
+  (e.g. branch a rollout from a saved state, or replay a fixed initial
+  condition), pass ``set_state=True``: ``env.reset(td, set_state=True)``. For
+  stateless environments such as :class:`~torchrl.envs.PendulumEnv` this honors
+  the state entries found in ``td``; for stateful environments that support it,
+  the underlying set-state API is used; envs that cannot honor a provided state
+  raise ``NotImplementedError``. ``set_state`` is a keyword argument (not a
+  tensordict key) so it never stacks/pads across a rollout. When ``set_state`` is
+  left unspecified but the tensordict carries state, the state is honored for
+  backwards compatibility and a :class:`FutureWarning` is emitted: from v0.15 an
+  unspecified ``set_state`` will be treated as ``False`` (state ignored, fresh
+  reset).
 - :meth:`env.step`: a step method that takes a :class:`tensordict.TensorDict` input
   containing an input action as well as other inputs (for model-based or stateless
   environments, for instance).

knowledge_base/ISAACLAB.md

@@ -85,11 +85,26 @@ print(td["policy"].shape)
 
 ## Tiled-camera rendering
 
-### Add a `TiledCameraCfg`
-
 Manager-based environments do not put camera data in the observation manager by
-default. Add a tiled camera to the scene config before `gym.make(...)` and read
-its batched RGB output.
+default. Add a tiled camera to the scene config before `gym.make(...)` and launch
+with `--enable_cameras`; without it, camera/rendering APIs are not initialized.
+
+For TorchRL, prefer `IsaacLabWrapper.add_tiled_camera_config(...)` and
+`IsaacLabWrapper(..., from_tiled_camera=True)` so pixels are inserted into the
+TensorDict under `"pixels"`.
+
+```python
+from torchrl.envs.libs.isaac_lab import IsaacLabWrapper
+
+IsaacLabWrapper.add_tiled_camera_config(env_cfg, width=64, height=64)
+env = IsaacLabWrapper(
+    gym.make("Isaac-Ant-v0", cfg=env_cfg),
+    device="cuda:0",
+    from_tiled_camera=True,
+)
+```
+
+If configuring the sensor manually, add a `TiledCameraCfg` to the scene config:
 
 ```python
 from isaaclab.sensors import TiledCameraCfg
@@ -114,16 +129,83 @@ env_cfg.scene.tiled_camera = TiledCameraCfg(
 )
 ```
 
-Use `--enable_cameras` with `AppLauncher`; without it, camera/rendering APIs are
-not initialized.
+`TiledCamera` renders all environments in a single batched pass on the GPU,
+producing `[num_envs, H, W, C]` tensors efficiently. Rendering many
+environments is expensive: start with a smaller number of envs for pixels, and
+increase `env_cfg.scene.env_spacing` if neighboring envs appear in camera views.
 
-For TorchRL, prefer `IsaacLabWrapper.add_tiled_camera_config(...)` and
-`IsaacLabWrapper(..., from_tiled_camera=True)` so pixels are inserted into the
-TensorDict under `"pixels"`.
+For the ANYmal-C quadruped (base at approximately 0.5 m height), useful camera
+positions include:
+
+- rear-elevated: `pos=(-3.0, 0.0, 2.0)`;
+- side view: `pos=(0.0, -3.0, 1.5)`;
+- top-down: `pos=(0.0, 0.0, 5.0)`.
+
+The rotation quaternion `(w, x, y, z) = (0.9945, 0.0, 0.1045, 0.0)` applies a
+slight downward pitch (approximately 12 degrees).
+
+## Auto-reset and per-index reset
+
+IsaacLab environments auto-reset individual sub-environments when they reach a
+terminal state. Done can be reported immediately after reset.
+
+`IsaacLabWrapper(env, native_autoreset=True)` keeps Isaac's native post-reset
+observation in `tensordict_["policy"]` and marks the terminal
+`("next", "policy")` with NaN; `EnvBase.step_and_maybe_reset` then skips the
+synthetic reset call. The same bridge is installed for Direct-workflow envs
+(`DirectRLEnv`, `DirectMARLEnv`), not just Manager-based envs.
+
+`IsaacLabWrapper` surfaces Isaac Lab's per-index reset and `reset_to` APIs
+through the standard torchrl `"_reset"` boolean mask:
+
+```python
+env = IsaacLabWrapper(gym.make("Isaac-Ant-v0", cfg=AntEnvCfg()), native_autoreset=True)
+td = env.reset()
+
+# ... step the env a few times ...
+
+# Reset half of the sub-envs without disturbing the others. The transform
+# stack (RewardSum, InitTracker, recurrent primers, VecNormV2, ...) fires
+# on the masked rows only, exactly like a normal reset.
+reset_mask = torch.zeros(env.batch_size[0], 1, dtype=torch.bool, device=env.device)
+reset_mask[: env.batch_size[0] // 2] = True
+td.set("_reset", reset_mask)
+env.reset(td)
+
+# Snapshot and branch from a deterministic state (manager-based envs only).
+snapshot = env.get_state()
+# ... evolve env from `snapshot` to explore one branch ...
+# Rewind to the snapshot via the unified deterministic-reset path:
+env.reset(td, set_state=True, scene_state=snapshot)
+
+# Convenience method (equivalent to the call above):
+env.reset_to_state(snapshot, td)
+```
 
-Rendering many environments is expensive. Start with a smaller number of envs
-for pixels, and increase `env_cfg.scene.env_spacing` if neighboring envs appear
-in camera views.
+Gotchas:
+
+- The per-index reset path is gated on `native_autoreset=True`. With the
+  default `native_autoreset=False`, the `VecGymEnvTransform`-based obs-swap path
+  already handles `step_and_maybe_reset`-driven partial resets implicitly; an
+  explicit per-index reset would double-reset those envs.
+- `reset_to_state` is only available on manager-based envs (`ManagerBasedEnv` /
+  `ManagerBasedRLEnv`). Direct envs do not expose `reset_to`.
+- `is_relative=True` interprets the snapshot pose relative to the env origin,
+  which is useful for terrain-relative pose reuse.
+- The snapshot is passed as the `scene_state` keyword argument rather than
+  inside the tensordict. A stateless env's reset state is torch-native
+  (tensordict / `torch.Tensor` entries in `state_spec`) and so lives in the
+  tensordict, but Isaac Lab's scene state is an opaque, non-torch-native object;
+  carrying it in the tensordict would require a `NonTensor` `state_spec` entry
+  threaded through the transform/step-MDP machinery, whereas a kwarg is simpler
+  and keeps simulator state out of the data path.
+
+## In-place tensor modification
+
+IsaacLab modifies the `terminated` and `truncated` tensors in-place.
+`IsaacLabWrapper` defensively clones `terminated`, `truncated` and `done` in its
+output transform, so the TensorDict returned by `step` is already safe from this
+aliasing; no extra handling is needed downstream.
 
 ### Headless EGL/Vulkan dependencies
 

test/envs/test_env_base.py

@@ -4,9 +4,11 @@
 # LICENSE file in the root directory of this source tree.
 from __future__ import annotations
 
+import argparse
 import functools
 import gc
 import pickle
+import warnings
 from functools import partial
 from typing import Any
 
@@ -20,7 +22,7 @@
 from tensordict.tensorclass import TensorClass
 from torch import nn
 
-from torchrl.data.tensor_specs import Composite, NonTensor, Unbounded
+from torchrl.data.tensor_specs import Binary, Composite, NonTensor, Unbounded
 from torchrl.envs import EnvBase, ParallelEnv, SerialEnv
 from torchrl.envs.libs.gym import gym_backend, GymEnv
 from torchrl.envs.transforms import StepCounter, TransformedEnv
@@ -918,3 +920,116 @@ def test_actions_and_policy_mutually_exclusive(self):
         env = CountingEnv(max_steps=10)
         with pytest.raises(ValueError, match="both"):
             env.rollout(3, policy=lambda td: td, actions=[torch.ones(1)])
+
+
+class _NestedStatelessEnv(EnvBase):
+    """A tiny stateless env whose state lives under a nested key.
+
+    Used to exercise the ``reset(td, set_state=True)`` path with a
+    :class:`~tensordict.NestedKey` state entry.
+    """
+
+    _supports_set_state = True
+
+    def __init__(self, **kwargs):
+        super().__init__(batch_size=(), **kwargs)
+        self.observation_spec = Composite(nested=Composite(x=Unbounded(shape=(1,))))
+        self.state_spec = self.observation_spec.clone()
+        self.action_spec = Unbounded(shape=(1,))
+        self.reward_spec = Unbounded(shape=(1,))
+        self.done_spec = Binary(n=1, shape=(1,), dtype=torch.bool)
+
+    def _reset(self, tensordict, **kwargs):
+        set_state = bool(kwargs.get("set_state"))
+        if (
+            set_state
+            and tensordict is not None
+            and ("nested", "x") in tensordict.keys(True)
+        ):
+            x = tensordict["nested", "x"].clone()
+        else:
+            x = torch.zeros(1)
+        return TensorDict(
+            {("nested", "x"): x, "done": torch.zeros(1, dtype=torch.bool)},
+            batch_size=(),
+        )
+
+    def _step(self, tensordict):
+        x = tensordict["nested", "x"] + tensordict["action"]
+        return TensorDict(
+            {
+                ("nested", "x"): x,
+                "reward": torch.zeros(1),
+                "done": torch.zeros(1, dtype=torch.bool),
+            },
+            batch_size=(),
+        )
+
+    def _set_seed(self, *args, **kwargs):
+        ...
+
+
+class TestResetSetState:
+    """Tests for the explicit ``reset(td, set_state=True)`` deterministic-reset kwarg."""
+
+    def test_set_state_nested_key(self):
+        env = _NestedStatelessEnv()
+        td = TensorDict({("nested", "x"): torch.full((1,), 7.0)}, batch_size=())
+        # honored when set_state=True
+        out = env.reset(td.clone(), set_state=True)
+        assert out["nested", "x"].item() == 7.0
+        # ignored when set_state=False
+        out_false = env.reset(td.clone(), set_state=False)
+        assert out_false["nested", "x"].item() == 0.0
+
+    def test_set_state_nested_key_transition_warning(self):
+        env = _NestedStatelessEnv()
+        td = TensorDict({("nested", "x"): torch.full((1,), 3.0)}, batch_size=())
+        # unspecified set_state with a populated nested state key -> FutureWarning,
+        # honored for backwards compatibility.
+        with pytest.warns(FutureWarning, match="set_state"):
+            out = env.reset(td.clone())
+        assert out["nested", "x"].item() == 3.0
+        # an empty tensordict must not warn.
+        with warnings.catch_warnings():
+            warnings.simplefilter("error", FutureWarning)
+            env.reset(TensorDict(batch_size=()))
+
+    def test_set_state_unsupported_raises(self):
+        # ContinuousActionVecMockEnv does not opt into deterministic resets.
+        env = ContinuousActionVecMockEnv()
+        assert env._supports_set_state is False
+        td = env.reset()
+        with pytest.raises(NotImplementedError, match="set_state"):
+            env.reset(td.clone(), set_state=True)
+        # an unsupported env with state in the td must not emit the warning.
+        with warnings.catch_warnings():
+            warnings.simplefilter("error", FutureWarning)
+            env.reset(td.clone())
+
+    def test_set_state_batched_serial(self):
+        env = SerialEnv(2, _NestedStatelessEnv)
+        try:
+            assert env._supports_set_state is True
+            td = env.reset()
+            td["nested", "x"] = torch.full((2, 1), 5.0)
+            out = env.reset(td.clone(), set_state=True)
+            assert (out["nested", "x"] == 5.0).all()
+        finally:
+            env.close()
+
+    def test_set_state_batched_parallel(self, maybe_fork_ParallelEnv):
+        env = maybe_fork_ParallelEnv(2, _NestedStatelessEnv)
+        try:
+            assert env._supports_set_state is True
+            td = env.reset()
+            td["nested", "x"] = torch.full((2, 1), 5.0)
+            out = env.reset(td.clone(), set_state=True)
+            assert (out["nested", "x"] == 5.0).all()
+        finally:
+            env.close()
+
+
+if __name__ == "__main__":
+    args, unknown = argparse.ArgumentParser().parse_known_args()
+    pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)