[BugFix] Fix IsaacLab reset regressions (#3869)

Author: vmoens

test/envs/test_env_base.py

@@ -25,7 +25,7 @@
 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
+from torchrl.envs.transforms import StepCounter, Transform, TransformedEnv
 from torchrl.envs.utils import check_env_specs, make_composite_from_td, step_mdp
 from torchrl.modules import Actor
 from torchrl.testing import (
@@ -969,6 +969,51 @@ def _set_seed(self, *args, **kwargs):
         ...
 
 
+class _KwargOnlySetStateEnv(EnvBase):
+    _supports_set_state = True
+
+    def __init__(self, **kwargs):
+        super().__init__(batch_size=(), **kwargs)
+        self.observation_spec = Composite(observation=Unbounded(shape=(1,)))
+        self.action_spec = Unbounded(shape=(1,))
+        self.reward_spec = Unbounded(shape=(1,))
+        self.done_spec = Binary(n=1, shape=(1,), dtype=torch.bool)
+
+    def _input_td_has_state(self, tensordict):
+        return False
+
+    def _reset(self, tensordict, **kwargs):
+        if kwargs.get("set_state"):
+            raise RuntimeError("unexpected implicit set_state")
+        return TensorDict(
+            {
+                "observation": torch.zeros(1),
+                "done": torch.zeros(1, dtype=torch.bool),
+            },
+            batch_size=(),
+        )
+
+    def _step(self, tensordict):
+        return TensorDict(
+            {
+                "observation": tensordict["observation"] + tensordict["action"],
+                "reward": torch.zeros(1),
+                "done": torch.zeros(1, dtype=torch.bool),
+            },
+            batch_size=(),
+        )
+
+    def _set_seed(self, *args, **kwargs):
+        ...
+
+
+class _OuterStateTransform(Transform):
+    def transform_state_spec(self, state_spec):
+        state_spec = state_spec.clone()
+        state_spec["outer_state"] = Unbounded(shape=(1,))
+        return state_spec
+
+
 class TestResetSetState:
     """Tests for the explicit ``reset(td, set_state=True)`` deterministic-reset kwarg."""
 
@@ -1029,6 +1074,14 @@ def test_set_state_batched_parallel(self, maybe_fork_ParallelEnv):
         finally:
             env.close()
 
+    def test_transformed_env_delegates_implicit_state_detection(self):
+        env = TransformedEnv(_KwargOnlySetStateEnv(), _OuterStateTransform())
+        td = TensorDict({"outer_state": torch.ones(1)}, batch_size=())
+        with warnings.catch_warnings():
+            warnings.simplefilter("error", FutureWarning)
+            out = env.reset(td)
+        assert (out["observation"] == 0).all()
+
 
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()

test/libs/test_isaac.py

@@ -9,15 +9,17 @@
 import itertools
 import os
 import queue as queue_lib
+import sys
 import time
 import traceback
+import types
 from functools import partial
 
 import pytest
 import torch
 import torch.distributed as dist
 import torchrl.testing.env_helper
-from tensordict import assert_allclose_td
+from tensordict import assert_allclose_td, TensorDict
 from tensordict.nn import TensorDictModule as Mod, TensorDictSequential as Seq
 from torch import multiprocessing as mp
 
@@ -28,6 +30,8 @@
 from torchrl.data.replay_buffers.samplers import SliceSampler
 from torchrl.data.replay_buffers.storages import LazyTensorStorage
 from torchrl.envs import InitTracker, RewardSum, StepCounter, TransformedEnv, VecNormV2
+from torchrl.envs.libs import gym as gym_lib, isaac_lab as isaac_lab_lib
+from torchrl.envs.libs.isaac_lab import IsaacLabWrapper
 from torchrl.envs.utils import check_env_specs
 from torchrl.modules import LSTMModule, MLP
 from torchrl.testing import get_default_devices
@@ -304,6 +308,105 @@ def _isaaclab_direct_native_autoreset(env_name: str, num_envs: int = 16):
         proc.join()
 
 
+def _install_fake_isaaclab(monkeypatch):
+    class ManagerBasedEnv:
+        pass
+
+    class DirectRLEnv:
+        pass
+
+    class DirectMARLEnv:
+        pass
+
+    fake_envs = types.ModuleType("isaaclab.envs")
+    fake_envs.ManagerBasedEnv = ManagerBasedEnv
+    fake_envs.DirectRLEnv = DirectRLEnv
+    fake_envs.DirectMARLEnv = DirectMARLEnv
+    fake_isaaclab = types.ModuleType("isaaclab")
+    fake_isaaclab.envs = fake_envs
+    monkeypatch.setitem(sys.modules, "isaaclab", fake_isaaclab)
+    monkeypatch.setitem(sys.modules, "isaaclab.envs", fake_envs)
+    return ManagerBasedEnv, DirectRLEnv, DirectMARLEnv
+
+
+def test_isaaclab_direct_env_detection_is_native_autoreset_opt_in(monkeypatch):
+    ManagerBasedEnv, DirectRLEnv, DirectMARLEnv = _install_fake_isaaclab(monkeypatch)
+    monkeypatch.setattr(gym_lib, "_has_isaaclab", True)
+    monkeypatch.setattr(isaac_lab_lib, "_has_isaaclab", True)
+
+    manager_env = ManagerBasedEnv()
+    direct_env = DirectRLEnv()
+    direct_marl_env = DirectMARLEnv()
+
+    assert IsaacLabWrapper._supports_native_autoreset(manager_env)
+    assert not IsaacLabWrapper._supports_native_autoreset(direct_env)
+    assert not IsaacLabWrapper._supports_native_autoreset(direct_marl_env)
+    assert IsaacLabWrapper._supports_native_autoreset(direct_env, native_autoreset=True)
+    assert IsaacLabWrapper._supports_native_autoreset(
+        direct_marl_env, native_autoreset=True
+    )
+
+    fake_vector = types.SimpleNamespace(VectorEnv=type("VectorEnv", (), {}))
+    monkeypatch.setattr(
+        gym_lib,
+        "gym_backend",
+        lambda name=None: fake_vector if name == "vector" else fake_vector,
+    )
+    wrapper = gym_lib.GymWrapper.__new__(gym_lib.GymWrapper)
+    wrapper._torchrl_native_autoreset_requested = False
+    wrapper._env = types.SimpleNamespace(unwrapped=manager_env)
+    assert wrapper._is_batched
+    wrapper._env = types.SimpleNamespace(unwrapped=direct_env)
+    assert not wrapper._is_batched
+    wrapper._torchrl_native_autoreset_requested = True
+    assert wrapper._is_batched
+
+    isaac_wrapper = IsaacLabWrapper.__new__(IsaacLabWrapper)
+    isaac_wrapper._env = types.SimpleNamespace(unwrapped=direct_env)
+    assert not isaac_wrapper._supports_set_state
+
+    def reset_to(*args, **kwargs):
+        return None
+
+    manager_env.reset_to = reset_to
+    isaac_wrapper._env = types.SimpleNamespace(unwrapped=manager_env)
+    assert isaac_wrapper._supports_set_state
+
+
+def test_isaaclab_observation_key_normalization_is_cached_and_non_clobbering():
+    env = IsaacLabWrapper.__new__(IsaacLabWrapper)
+    env._rename_policy_to_observation = False
+    policy = torch.ones(2, 3)
+    observations = {"policy": policy}
+    assert env._normalize_observation_keys(observations) is observations
+
+    env._rename_policy_to_observation = True
+    normalized = env._normalize_observation_keys(observations)
+    assert normalized is not observations
+    assert "policy" not in normalized
+    assert normalized["observation"] is policy
+
+    existing_observation = torch.zeros(2, 3)
+    observations = {"policy": policy, "observation": existing_observation}
+    assert env._normalize_observation_keys(observations) is observations
+    assert observations["observation"] is existing_observation
+
+
+def test_isaaclab_all_false_reset_to_state_is_no_op():
+    env = IsaacLabWrapper.__new__(IsaacLabWrapper)
+    td = TensorDict(
+        {
+            "_reset": torch.zeros(3, 1, dtype=torch.bool),
+            "policy": torch.ones(3, 2),
+        },
+        batch_size=(3,),
+    )
+    out = env._reset(td, set_state=True, scene_state=object())
+    assert "_reset" not in out.keys()
+    assert out is not td
+    assert (out["policy"] == td["policy"]).all()
+
+
 @pytest.mark.skipif(not _has_isaac, reason="IsaacGym not found")
 @pytest.mark.parametrize(
     "task",

torchrl/envs/libs/gym.py

@@ -873,6 +873,9 @@ def __call__(cls, *args, **kwargs):
         missing_obs_value = kwargs.pop("missing_obs_value", None)
         native_autoreset = kwargs.pop("native_autoreset", False)
         num_workers = kwargs.pop("num_workers", 1)
+        native_autoreset = kwargs.setdefault(
+            "_torchrl_native_autoreset_requested", native_autoreset
+        )
 
         if cls.__name__ == "GymEnv" and num_workers > 1:
             from torchrl.envs import EnvCreator, ParallelEnv
@@ -903,7 +906,9 @@ def __call__(cls, *args, **kwargs):
                 kwargs = {}
                 if missing_obs_value is not None:
                     kwargs["missing_obs_value"] = missing_obs_value
-                if IsaacLabWrapper._supports_native_autoreset(instance._env.unwrapped):
+                if IsaacLabWrapper._supports_native_autoreset(
+                    instance._env.unwrapped, native_autoreset=native_autoreset
+                ):
                     env = TransformedEnv(
                         instance,
                         VecGymEnvTransform(**kwargs, native_autoreset=native_autoreset),
@@ -1131,6 +1136,9 @@ def get_library_name(env) -> str:
         )
 
     def __init__(self, env=None, categorical_action_encoding=False, **kwargs):
+        self._torchrl_native_autoreset_requested = kwargs.pop(
+            "_torchrl_native_autoreset_requested", False
+        )
         self._seed_calls_reset = None
         self._categorical_action_encoding = categorical_action_encoding
         if env is not None:
@@ -1204,7 +1212,10 @@ def _is_batched(self):
             from torchrl.envs.libs.isaac_lab import IsaacLabWrapper
 
             tuple_of_classes = (
-                tuple_of_classes + IsaacLabWrapper._supported_isaac_env_classes()
+                tuple_of_classes
+                + IsaacLabWrapper._supported_isaac_env_classes(
+                    include_direct=self._torchrl_native_autoreset_requested
+                )
             )
         return isinstance(
             self._env.unwrapped, tuple_of_classes + (gym_backend("vector").VectorEnv,)