[Task Clean-up] Dexterous Part 4/9: Fix MARL-to-single-agent training and enable handover Direct RSL-RL

docs/source/overview/environments.rst

@@ -387,7 +387,7 @@ for the lift-cube environment:
 .. |cube-shadow-link| replace:: :isaaclab-source:`Isaac-Reorient-Cube-Shadow-Direct <source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_env_cfg.py>`
 .. |cube-shadow-ff-link| replace:: :isaaclab-source:`Isaac-Reorient-Cube-Shadow-OpenAI-FF-Direct <source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_env_cfg.py>`
 .. |cube-shadow-lstm-link| replace:: :isaaclab-source:`Isaac-Reorient-Cube-Shadow-OpenAI-LSTM-Direct <source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_env_cfg.py>`
-.. |cube-shadow-vis-link| replace:: :isaaclab-source:`Isaac-Reorient-Cube-Shadow-Camera-Direct <source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_camera_env.py>`
+.. |cube-shadow-vis-link| replace:: :isaaclab-source:`Isaac-Reorient-Cube-Shadow-Camera-Direct <source/isaaclab_tasks/isaaclab_tasks/core/reorient/config/shadow_hand/shadow_hand_direct_camera_env.py>`
 .. |agibot_place_mug-link| replace:: :isaaclab-source:`IsaacContrib-Place-Mug-Agibot-Left-Arm-RmpFlow <source/isaaclab_tasks/isaaclab_tasks/contrib/place/config/agibot/place_upright_mug_rmp_rel_env_cfg.py>`
 .. |agibot_place_toy-link| replace:: :isaaclab-source:`IsaacContrib-Place-Toy2Box-Agibot-Right-Arm-RmpFlow <source/isaaclab_tasks/isaaclab_tasks/contrib/place/config/agibot/place_toy2box_rmp_rel_env_cfg.py>`
 .. |reach_openarm_bi-link| replace:: :isaaclab-source:`IsaacContrib-Reach-OpenArmBi <source/isaaclab_tasks/isaaclab_tasks/contrib/reach/config/openarm/bimanual/joint_pos_env_cfg.py>`

source/isaaclab/changelog.d/dexterous-env-convergence.rst

@@ -0,0 +1,6 @@
+Fixed
+^^^^^
+
+* Fixed :meth:`~isaaclab.envs.DirectRLEnv.reset` to store the observation buffer like
+  :meth:`~isaaclab.envs.DirectRLEnv.step`, and exposed it on the
+  multi-agent-to-single-agent adapter.

source/isaaclab/isaaclab/envs/direct_rl_env.py

@@ -386,7 +386,9 @@ def reset(self, seed: int | None = None, options: dict[str, Any] | None = None)
                     self.sim.render()
 
         # return observations
-        return self._get_observations(), self.extras
+        # store the buffer like step() does, so consumers can read the latest observations
+        self.obs_buf = self._get_observations()
+        return self.obs_buf, self.extras
 
     def step(self, action: torch.Tensor) -> VecEnvStepReturn:
         """Execute one time-step of the environment's dynamics.

source/isaaclab/isaaclab/envs/utils/marl.py

@@ -81,22 +81,39 @@ def __init__(self, env: DirectMARLEnv) -> None:
             )
             self.action_space = gym.vector.utils.batch_space(self.single_action_space, self.num_envs)
 
-        def reset(self, seed: int | None = None, options: dict[str, Any] | None = None) -> tuple[VecEnvObs, dict]:
-            obs, extras = self.env.reset(seed, options)
+            # latest converted observations, refreshed by reset() and step()
+            self._obs_buf: VecEnvObs = {}
 
-            # use environment state as observation
-            if self._state_as_observation:
-                obs = {"policy": self.env.state()}
-            # concatenate agents' observations
+        @property
+        def episode_length_buf(self) -> torch.Tensor:
+            """Episode lengths from the wrapped multi-agent environment."""
+            return self.env.episode_length_buf
+
+        @episode_length_buf.setter
+        def episode_length_buf(self, value: torch.Tensor) -> None:
+            # copy in place so holders of the wrapped environment's buffer stay in sync
+            self.env.episode_length_buf.copy_(value)
+
+        @property
+        def obs_buf(self) -> VecEnvObs:
+            """Latest observations from the wrapped multi-agent environment."""
+            return self._obs_buf
+
+        def _convert_observations(self, obs: dict[AgentID, ObsType]) -> VecEnvObs:
+            """Convert multi-agent observations to the single-agent policy observation."""
             # FIXME: This implementation assumes the spaces are fundamental ones. Fix it to support composite spaces
-            else:
-                obs = {
-                    "policy": torch.cat(
-                        [obs[agent].reshape(self.num_envs, -1) for agent in self.env.possible_agents], dim=-1
-                    )
-                }
+            if self._state_as_observation:
+                return {"policy": self.env.state()}
+            return {
+                "policy": torch.cat(
+                    [obs[agent].reshape(self.num_envs, -1) for agent in self.env.possible_agents], dim=-1
+                )
+            }
 
-            return obs, extras
+        def reset(self, seed: int | None = None, options: dict[str, Any] | None = None) -> tuple[VecEnvObs, dict]:
+            obs, extras = self.env.reset(seed, options)
+            self._obs_buf = self._convert_observations(obs)
+            return self._obs_buf, extras
 
         def step(self, action: torch.Tensor) -> VecEnvStepReturn:
             # split single-agent actions to build the multi-agent ones
@@ -111,24 +128,14 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn:
             # step the environment
             obs, rewards, terminated, time_outs, extras = self.env.step(_actions)
 
-            # use environment state as observation
-            if self._state_as_observation:
-                obs = {"policy": self.env.state()}
-            # concatenate agents' observations
-            # FIXME: This implementation assumes the spaces are fundamental ones. Fix it to support composite spaces
-            else:
-                obs = {
-                    "policy": torch.cat(
-                        [obs[agent].reshape(self.num_envs, -1) for agent in self.env.possible_agents], dim=-1
-                    )
-                }
+            self._obs_buf = self._convert_observations(obs)
 
             # process environment outputs to return single-agent data
             rewards = sum(rewards.values())
             terminated = math.prod(terminated.values()).to(dtype=torch.bool)
             time_outs = math.prod(time_outs.values()).to(dtype=torch.bool)
 
-            return obs, rewards, terminated, time_outs, extras
+            return self._obs_buf, rewards, terminated, time_outs, extras
 
         def render(self, recompute: bool = False) -> np.ndarray | None:
             return self.env.render(recompute)

source/isaaclab/test/envs/test_marl_utils.py

@@ -0,0 +1,109 @@
+# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+from types import SimpleNamespace
+
+import gymnasium as gym
+import torch
+
+from isaaclab.envs.utils.marl import multi_agent_to_single_agent
+
+
+class _FakeMultiAgentEnv:
+    possible_agents = ["agent_0", "agent_1"]
+    observation_spaces = {
+        "agent_0": gym.spaces.Box(low=-1.0, high=1.0, shape=(2,)),
+        "agent_1": gym.spaces.Box(low=-1.0, high=1.0, shape=(1,)),
+    }
+    action_spaces = {
+        "agent_0": gym.spaces.Box(low=-1.0, high=1.0, shape=(1,)),
+        "agent_1": gym.spaces.Box(low=-1.0, high=1.0, shape=(1,)),
+    }
+    render_mode = None
+
+    def __init__(self):
+        self.unwrapped = self
+        self.cfg = SimpleNamespace(state_space=2)
+        self.state_space = gym.spaces.Box(low=-1.0, high=1.0, shape=(2,))
+        self.sim = object()
+        self.scene = SimpleNamespace(num_envs=2)
+        self.episode_length_buf = torch.tensor([1, 2])
+        self.obs_dict = {
+            "agent_0": torch.tensor([[1.0, 2.0], [3.0, 4.0]]),
+            "agent_1": torch.tensor([[5.0], [6.0]]),
+        }
+
+    def reset(self, seed=None, options=None):
+        return self.obs_dict, {}
+
+    def step(self, actions):
+        # shift the observations so a step is distinguishable from a reset
+        self.obs_dict = {agent: obs + 10.0 for agent, obs in self.obs_dict.items()}
+        rewards = {agent: torch.zeros(2) for agent in self.possible_agents}
+        dones = {agent: torch.zeros(2, dtype=torch.bool) for agent in self.possible_agents}
+        return self.obs_dict, rewards, dones, dones, {}
+
+    def state(self):
+        return torch.tensor([[7.0, 8.0], [9.0, 10.0]])
+
+    def close(self):
+        pass
+
+
+def test_multi_agent_to_single_agent_reset_concatenates_agents():
+    """The adapter reset should concatenate the agents' observations."""
+    env = multi_agent_to_single_agent(_FakeMultiAgentEnv())
+
+    observations, _ = env.reset()
+
+    torch.testing.assert_close(observations["policy"], torch.tensor([[1.0, 2.0, 5.0], [3.0, 4.0, 6.0]]))
+
+
+def test_multi_agent_to_single_agent_reset_can_use_state():
+    """The adapter reset should support the state-as-observation mode."""
+    env = multi_agent_to_single_agent(_FakeMultiAgentEnv(), state_as_observation=True)
+
+    observations, _ = env.reset()
+
+    torch.testing.assert_close(observations["policy"], torch.tensor([[7.0, 8.0], [9.0, 10.0]]))
+
+
+def test_multi_agent_to_single_agent_forwards_episode_lengths():
+    """RSL-RL episode randomization should update the wrapped environment buffer in place."""
+    source_env = _FakeMultiAgentEnv()
+    wrapped_buffer = source_env.episode_length_buf
+    env = multi_agent_to_single_agent(source_env)
+    episode_lengths = torch.tensor([3, 4])
+
+    env.episode_length_buf = episode_lengths
+
+    torch.testing.assert_close(env.episode_length_buf, episode_lengths)
+    torch.testing.assert_close(source_env.episode_length_buf, episode_lengths)
+    # written in place, so references taken before the assignment observe the new values
+    assert source_env.episode_length_buf is wrapped_buffer
+
+
+def test_multi_agent_to_single_agent_exposes_latest_observations():
+    """The public observation buffer should track the observations last returned by reset and step."""
+    env = multi_agent_to_single_agent(_FakeMultiAgentEnv())
+
+    reset_obs, _ = env.reset()
+    torch.testing.assert_close(env.obs_buf["policy"], reset_obs["policy"])
+    torch.testing.assert_close(env.obs_buf["policy"], torch.tensor([[1.0, 2.0, 5.0], [3.0, 4.0, 6.0]]))
+
+    step_obs = env.step(torch.zeros(2, 2))[0]
+    torch.testing.assert_close(env.obs_buf["policy"], step_obs["policy"])
+    torch.testing.assert_close(env.obs_buf["policy"], torch.tensor([[11.0, 12.0, 15.0], [13.0, 14.0, 16.0]]))
+
+
+def test_multi_agent_to_single_agent_state_observation_tracks_steps():
+    """The state-as-observation mode should also refresh the buffer on every transition."""
+    env = multi_agent_to_single_agent(_FakeMultiAgentEnv(), state_as_observation=True)
+
+    reset_obs, _ = env.reset()
+    torch.testing.assert_close(env.obs_buf["policy"], reset_obs["policy"])
+
+    step_obs = env.step(torch.zeros(2, 2))[0]
+    torch.testing.assert_close(env.obs_buf["policy"], step_obs["policy"])

source/isaaclab_assets/changelog.d/task-cleanup-dex-part04.minor.rst

@@ -0,0 +1,5 @@
+Added
+^^^^^
+
+* Added ``SHADOW_HAND_NEWTON_CFG``, the Newton (MJWarp) Shadow Hand configuration,
+  shared by the reorientation and handover tasks.

source/isaaclab_assets/isaaclab_assets/robots/shadow_hand.py

@@ -7,7 +7,8 @@
 
 The following configurations are available:
 
-* :obj:`SHADOW_HAND_CFG`: Shadow Hand with implicit actuator model.
+* :obj:`SHADOW_HAND_CFG`: Shadow Hand on the PhysX asset with implicit actuator model.
+* :obj:`SHADOW_HAND_NEWTON_CFG`: Shadow Hand on the Newton (MJWarp) asset.
 
 Reference:
 
@@ -81,7 +82,97 @@
     },
     soft_joint_pos_limit_factor=1.0,
 )
-"""Configuration of Shadow Hand robot."""
+"""Configuration of the Shadow Hand robot on the PhysX asset."""
+
+
+SHADOW_HAND_NEWTON_CFG = ArticulationCfg(
+    spawn=sim_utils.UsdFileCfg(
+        # Newton/MuJoCo use a separate USD schema; this asset renumbers the finger
+        # joints (+1) relative to the PhysX asset above (e.g. FFJ4 vs FFJ3, LFJ5 vs LFJ4).
+        usd_path=f"{ISAAC_NUCLEUS_DIR}/Robots/ShadowRobot/ShadowHandNewton/shadow_hand_instanceable.usda",
+        activate_contact_sensors=False,
+        rigid_props=sim_utils.RigidBodyPropertiesCfg(
+            disable_gravity=True,
+            retain_accelerations=True,
+            max_depenetration_velocity=1000.0,
+        ),
+        articulation_props=sim_utils.ArticulationRootPropertiesCfg(enabled_self_collisions=True),
+        joint_drive_props=sim_utils.JointDrivePropertiesCfg(drive_type="force", ensure_drives_exist=True),
+        fixed_tendons_props=sim_utils.FixedTendonPropertiesCfg(damping=0.1),
+    ),
+    init_state=ArticulationCfg.InitialStateCfg(
+        pos=(0.0, 0.0, 0.5),
+        # WARNING(Octi): Newton's import_usd.py bakes the USD body xformOp rotation into
+        # joint_X_p for the root fixed joint, which cancels with the matching localPose1
+        # rotation in joint_X_c during FK (joint_X_p * inv(joint_X_c) ≈ identity). This
+        # discards the root body's native USD orientation, so we must re-apply it here as a
+        # spawn rotation. PhysX or USD does not have this issue. Remove once Newton fixes root joint
+        # transform handling in import_usd.py.
+        rot=(0.0, 0.0, -0.70710678118, 0.70710678118),
+        joint_pos={".*": 0.0},
+    ),
+    actuators={
+        # Drives the joints named by :obj:`SHADOW_ACTUATED_JOINT_NAMES`, which resolve on this
+        # asset despite its +1 finger renumbering.
+        #
+        # The per-finger ``J1``/``J2`` pair is coupled by a fixed tendon (``coef=[1, 1]``) that
+        # the MJWarp solver currently skips, so the configuration is what holds the pair
+        # together: both ends are driven with identical gains and effort limits, which
+        # reproduces the 1:1 coupling. Keep them symmetric and in the same actuator group --
+        # driving one end while clamping the other makes the two fight, and an uncapped effort
+        # limit on either end diverges to NaN within a few hundred steps.
+        #
+        # Known limitation: the ``J4`` knuckle abduction joints (and ``LFJ5``) are left
+        # undriven, so the fingers cannot spread laterally. Correcting that requires resolving
+        # the skipped tendon first and is deferred to a follow-up.
+        "fingers": ImplicitActuatorCfg(
+            joint_names_expr=[
+                "robot0_WR.*",
+                "robot0_(FF|MF|RF|LF|TH)J(3|2|1)",
+                "robot0_(LF|TH)J4",
+                "robot0_THJ0",
+            ],
+            effort_limit_sim={
+                "robot0_WRJ1": 4.785,
+                "robot0_WRJ0": 2.175,
+                "robot0_(FF|MF|RF|LF)J1": 0.7245,
+                "robot0_FFJ(3|2)": 0.9,
+                "robot0_MFJ(3|2)": 0.9,
+                "robot0_RFJ(3|2)": 0.9,
+                "robot0_LFJ(4|3|2)": 0.9,
+                "robot0_THJ4": 2.3722,
+                "robot0_THJ3": 1.45,
+                "robot0_THJ(2|1)": 0.99,
+                "robot0_THJ0": 0.81,
+            },
+            # Default gains match the PhysX cfg (wrists 5.0/0.5, fingers 1.0/0.1). Tasks that
+            # need more joint authority override these -- e.g. the handover catch on MJWarp
+            # raises them to 20.0/2.0, since MJWarp's implicit-PD path lacks PhysX's
+            # fixed-tendon limit stiffness + solver-iteration torque amplification.
+            stiffness={
+                "robot0_WRJ.*": 5.0,
+                "robot0_(FF|MF|RF|LF|TH)J(3|2|1)": 1.0,
+                "robot0_(LF|TH)J4": 1.0,
+                "robot0_THJ0": 1.0,
+            },
+            damping={
+                "robot0_WRJ.*": 0.5,
+                "robot0_(FF|MF|RF|LF|TH)J(3|2|1)": 0.1,
+                "robot0_(LF|TH)J4": 0.1,
+                "robot0_THJ0": 0.1,
+            },
+            friction=1e-2,
+            armature=2e-3,
+        ),
+    },
+    soft_joint_pos_limit_factor=1.0,
+)
+"""Configuration of the Shadow Hand robot on the Newton (MJWarp) asset.
+
+The Newton USD renumbers the finger joints (+1) relative to :obj:`SHADOW_HAND_CFG`, but the
+names in :obj:`SHADOW_ACTUATED_JOINT_NAMES` resolve on both assets, so the two backends share
+one actuated-joint list. Gains default to the PhysX values; tasks override them as needed.
+"""
 
 
 SHADOW_FINGERTIP_BODY_NAMES: list[str] = [
@@ -115,4 +206,7 @@
     "robot0_THJ1",
     "robot0_THJ0",
 ]
-"""Shadow Hand actuated joint names, in the Direct task's actuation order."""
+"""Shadow Hand actuated joint names, in the Direct task's actuation order.
+
+These names resolve on both the PhysX and Newton assets, so every backend shares this list.
+"""

source/isaaclab_experimental/changelog.d/dexterous-env-convergence.rst

@@ -0,0 +1,6 @@
+Fixed
+^^^^^
+
+* Fixed :meth:`~isaaclab_experimental.envs.DirectRLEnvWarp.step` and
+  :meth:`~isaaclab_experimental.envs.DirectRLEnvWarp.reset` to store observations in
+  ``obs_buf``.

source/isaaclab_experimental/isaaclab_experimental/envs/direct_rl_env_warp.py

@@ -372,7 +372,9 @@ def reset(self, seed: int | None = None, options: dict[str, Any] | None = None)
 
         # return observations
         self._get_observations()
-        return {"policy": self.torch_obs_buf.clone()}, self.extras
+        # store the returned buffer so RslRlVecEnvWrapper.get_observations() can read env.obs_buf
+        self.obs_buf = {"policy": self.torch_obs_buf.clone()}
+        return self.obs_buf, self.extras
 
     @Timer(name="env_step", msg="Step took:", enable=DEBUG_TIMER_STEP or DEBUG_TIMERS)
     def step(self, action: torch.Tensor) -> VecEnvStepReturn:
@@ -455,8 +457,10 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn:
             self._post_step_visualize()
 
         # return observations, rewards, resets and extras
+        # store the returned buffer so RslRlVecEnvWrapper.get_observations() can read env.obs_buf
+        self.obs_buf = {"policy": self.torch_obs_buf.clone()}
         return (
-            {"policy": self.torch_obs_buf.clone()},
+            self.obs_buf,
             self.torch_reward_buf,
             self.torch_reset_terminated,
             self.torch_reset_time_outs,

source/isaaclab_rl/changelog.d/dexterous-env-convergence.minor.rst

@@ -0,0 +1,7 @@
+Changed
+^^^^^^^
+
+* Changed :meth:`~isaaclab_rl.rsl_rl.RslRlVecEnvWrapper.get_observations` to read the
+  environment-owned observation buffer instead of private environment methods. The
+  returned observations now match the latest reset/step returns, including observation
+  noise.

source/isaaclab_rl/isaaclab_rl/rsl_rl/vecenv_wrapper.py

@@ -171,11 +171,7 @@ def reset(self) -> tuple[TensorDict, dict]:  # noqa: D102
 
     def get_observations(self) -> TensorDict:
         """Returns the current observations of the environment."""
-        if hasattr(self.unwrapped, "observation_manager"):
-            obs_dict = self.unwrapped.observation_manager.compute()
-        else:
-            obs_dict = self.unwrapped._get_observations()
-        return TensorDict(obs_dict, batch_size=[self.num_envs])
+        return TensorDict(self.unwrapped.obs_buf, batch_size=[self.num_envs])
 
     def step(self, actions: torch.Tensor) -> tuple[TensorDict, torch.Tensor, torch.Tensor, dict]:
         # clip actions