From 6cf3e87ebc19ae54c39c5ea1f346f944ec2b5deb Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Sun, 30 Aug 2026 14:14:58 -0700 Subject: [PATCH 1/4] Make zero agent use neutral actions --- .../mhaiderbhai-neutral-actions.minor.rst | 5 ++ .../mdp/actions/pink_task_space_actions.py | 23 ++++++ .../mdp/actions/rmpflow_task_space_actions.py | 10 +++ .../envs/mdp/actions/task_space_actions.py | 44 +++++++++++ .../isaaclab/managers/action_manager.py | 21 ++++++ .../test/envs/test_neutral_actions.py | 43 +++++++++++ .../mhaiderbhai-zero-agent-actions.rst | 5 ++ .../isaaclab_rl/entrypoints/simple_agents.py | 47 +++++++++--- source/isaaclab_rl/test/test_entrypoints.py | 75 +++++++++++++++++++ 9 files changed, 263 insertions(+), 10 deletions(-) create mode 100644 source/isaaclab/changelog.d/mhaiderbhai-neutral-actions.minor.rst create mode 100644 source/isaaclab/test/envs/test_neutral_actions.py create mode 100644 source/isaaclab_rl/changelog.d/mhaiderbhai-zero-agent-actions.rst diff --git a/source/isaaclab/changelog.d/mhaiderbhai-neutral-actions.minor.rst b/source/isaaclab/changelog.d/mhaiderbhai-neutral-actions.minor.rst new file mode 100644 index 000000000000..9343ebf809c6 --- /dev/null +++ b/source/isaaclab/changelog.d/mhaiderbhai-neutral-actions.minor.rst @@ -0,0 +1,5 @@ +Added +^^^^^ + +* Added semantic neutral actions for manager-based action terms, including absolute differential IK, + Pink IK, RMPFlow, and operational-space controllers. diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py index 3a5e67ffc426..32949738f918 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py @@ -75,6 +75,16 @@ def _initialize_joint_info(self) -> None: # Resolve hand joints self._hand_joint_ids, self._hand_joint_names = self._asset.find_joints(self.cfg.hand_joint_names) + # Resolve controlled frames in the same order as their pose commands. + self._controlled_frame_ids, controlled_frame_names = self._asset.find_bodies( + list(self.cfg.target_eef_link_names.values()), preserve_order=True + ) + if len(self._controlled_frame_ids) != len(self.cfg.target_eef_link_names): + raise ValueError( + "Expected one controlled body for every Pink IK target. Resolved " + f"{controlled_frame_names} from {list(self.cfg.target_eef_link_names.values())}." + ) + # Combine all joint information self._controlled_joint_ids = self._isaaclab_controlled_joint_ids + self._hand_joint_ids self._controlled_joint_names = self._isaaclab_controlled_joint_names + self._hand_joint_names @@ -109,6 +119,11 @@ def _initialize_helper_tensors(self) -> None: 1 for task in self._ik_controllers[0].cfg.variable_input_tasks if isinstance(task, FrameTask) ) self._num_frame_tasks = num_frame_tasks + if len(self._controlled_frame_ids) != self._num_frame_tasks: + raise ValueError( + f"Pink IK has {self._num_frame_tasks} variable frame tasks but " + f"{len(self._controlled_frame_ids)} controlled bodies were configured." + ) self._controlled_frame_poses = torch.zeros(num_frame_tasks, self.num_envs, 4, 4, device=self.device) # Pre-allocate tensor for base frame computations @@ -155,6 +170,14 @@ def processed_actions(self) -> torch.Tensor: """Get the processed actions tensor.""" return self._processed_actions + @property + def neutral_actions(self) -> torch.Tensor: + """Raw actions that hold the controlled frames and hand joints at their current state.""" + frame_poses = self._asset.data.body_link_pose_w.torch[:, self._controlled_frame_ids].clone() + frame_poses[..., :3] -= self._env.scene.env_origins.unsqueeze(1) + hand_joint_positions = self._asset.data.joint_pos.torch[:, self._hand_joint_ids] + return torch.cat((frame_poses.flatten(start_dim=1), hand_joint_positions), dim=-1) + @property def IO_descriptor(self) -> GenericActionIODescriptor: """The IO descriptor of the action term. diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/rmpflow_task_space_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/rmpflow_task_space_actions.py index 6cccf308e1d5..1893cc465f78 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/rmpflow_task_space_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/rmpflow_task_space_actions.py @@ -118,6 +118,16 @@ def raw_actions(self) -> torch.Tensor: def processed_actions(self) -> torch.Tensor: return self._processed_actions + @property + def neutral_actions(self) -> torch.Tensor: + """Raw actions that hold the current end-effector pose.""" + if self.cfg.use_relative_mode: + return super().neutral_actions + + ee_pos, ee_quat = self._compute_frame_pose() + command = torch.cat((ee_pos, ee_quat), dim=-1) + return torch.where(self._scale != 0.0, command / self._scale, torch.zeros_like(command)) + @property def jacobian_w(self) -> torch.Tensor: return self._asset.data.body_link_jacobian_w.torch[:, self._jacobi_body_idx, :, self._jacobi_joint_ids] diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py index ddc994ec52b0..33adfba89b32 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py @@ -139,6 +139,16 @@ def raw_actions(self) -> torch.Tensor: def processed_actions(self) -> torch.Tensor: return self._processed_actions + @property + def neutral_actions(self) -> torch.Tensor: + """Raw actions that hold the current end-effector pose.""" + if self.cfg.controller.use_relative_mode: + return super().neutral_actions + + ee_pos, ee_quat = self._compute_frame_pose() + command = ee_pos if self.cfg.controller.command_type == "position" else torch.cat((ee_pos, ee_quat), dim=-1) + return torch.where(self._scale != 0.0, command / self._scale, torch.zeros_like(command)) + @property def jacobian_w(self) -> torch.Tensor: return self._asset.data.body_link_jacobian_w.torch[:, self._jacobi_body_idx, :, self._jacobi_joint_ids] @@ -450,6 +460,40 @@ def processed_actions(self) -> torch.Tensor: """Processed actions for operational space control.""" return self._processed_actions + @property + def neutral_actions(self) -> torch.Tensor: + """Raw actions that hold the current end-effector pose and apply no wrench.""" + actions = super().neutral_actions + if self._pose_abs_idx is None: + return actions + + self._compute_ee_pose() + self._compute_task_frame_pose() + if self._task_frame_pose_b is None: + ee_pos_task = self._ee_pose_b[:, :3] + ee_quat_task = self._ee_pose_b[:, 3:7] + else: + ee_pos_task, ee_quat_task = math_utils.subtract_frame_transforms( + self._task_frame_pose_b[:, :3], + self._task_frame_pose_b[:, 3:7], + self._ee_pose_b[:, :3], + self._ee_pose_b[:, 3:7], + ) + + position_slice = slice(self._pose_abs_idx, self._pose_abs_idx + 3) + orientation_slice = slice(self._pose_abs_idx + 3, self._pose_abs_idx + 7) + actions[:, position_slice] = torch.where( + self._position_scale != 0.0, + ee_pos_task / self._position_scale, + torch.zeros_like(ee_pos_task), + ) + actions[:, orientation_slice] = torch.where( + self._orientation_scale != 0.0, + ee_quat_task / self._orientation_scale, + torch.zeros_like(ee_quat_task), + ) + return actions + @property def jacobian_w(self) -> torch.Tensor: return self._asset.data.body_link_jacobian_w.torch[:, self._jacobi_ee_body_idx, :, self._jacobi_joint_idx] diff --git a/source/isaaclab/isaaclab/managers/action_manager.py b/source/isaaclab/isaaclab/managers/action_manager.py index d711596e5f5a..1e952628b168 100644 --- a/source/isaaclab/isaaclab/managers/action_manager.py +++ b/source/isaaclab/isaaclab/managers/action_manager.py @@ -88,6 +88,15 @@ def processed_actions(self) -> torch.Tensor: """The actions computed by the term after applying any processing.""" raise NotImplementedError + @property + def neutral_actions(self) -> torch.Tensor: + """Raw actions suitable for passive agent playback. + + The default is a zero-filled tensor. Action terms for which zero has a different or invalid meaning, + such as absolute-pose controllers, should override this property with a semantically neutral command. + """ + return torch.zeros_like(self.raw_actions) + @property def has_debug_vis_implementation(self) -> bool: """Whether the action term has a debug visualization implemented.""" @@ -263,6 +272,18 @@ def prev_action(self) -> torch.Tensor: """The previous actions sent to the environment. Shape is (num_envs, total_action_dim).""" return self._prev_action + @property + def neutral_actions(self) -> torch.Tensor: + """Raw actions suitable for passive playback of all active action terms. + + The returned tensor has shape ``(num_envs, total_action_dim)``. Since + some terms derive their neutral command from the current simulation + state, consumers should retrieve this property immediately before use. + """ + if not self._terms: + return torch.zeros_like(self._action) + return torch.cat([term.neutral_actions for term in self._terms.values()], dim=-1) + @property def has_debug_vis_implementation(self) -> bool: """Whether the command terms have debug visualization implemented.""" diff --git a/source/isaaclab/test/envs/test_neutral_actions.py b/source/isaaclab/test/envs/test_neutral_actions.py new file mode 100644 index 000000000000..d0b6708c314d --- /dev/null +++ b/source/isaaclab/test/envs/test_neutral_actions.py @@ -0,0 +1,43 @@ +# 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 + +"""Tests for semantic neutral actions.""" + +from types import SimpleNamespace + +import torch + +from isaaclab.envs.mdp.actions.pink_task_space_actions import PinkInverseKinematicsAction + + +def test_pink_neutral_actions_use_current_frame_poses() -> None: + """Pink IK neutral actions contain valid current poses and hand joint positions.""" + action_term = object.__new__(PinkInverseKinematicsAction) + action_term._controlled_frame_ids = [1, 0] + action_term._hand_joint_ids = [1, 3] + + body_poses = torch.tensor( + [ + [[1.0, 2.0, 3.0, 0.0, 0.0, 0.0, 1.0], [4.0, 5.0, 6.0, 0.0, 0.0, 1.0, 0.0]], + [[7.0, 8.0, 9.0, 0.0, 1.0, 0.0, 0.0], [10.0, 11.0, 12.0, 1.0, 0.0, 0.0, 0.0]], + ] + ) + joint_positions = torch.tensor([[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8]]) + env_origins = torch.tensor([[1.0, 1.0, 1.0], [2.0, 2.0, 2.0]]) + action_term._asset = SimpleNamespace( + data=SimpleNamespace( + body_link_pose_w=SimpleNamespace(torch=body_poses), + joint_pos=SimpleNamespace(torch=joint_positions), + ) + ) + action_term._env = SimpleNamespace(scene=SimpleNamespace(env_origins=env_origins)) + + actions = action_term.neutral_actions + + expected_poses = body_poses[:, [1, 0]].clone() + expected_poses[..., :3] -= env_origins.unsqueeze(1) + expected = torch.cat((expected_poses.flatten(start_dim=1), joint_positions[:, [1, 3]]), dim=-1) + assert torch.equal(actions, expected) + assert torch.all(torch.linalg.vector_norm(actions[:, :14].reshape(2, 2, 7)[..., 3:7], dim=-1) == 1.0) diff --git a/source/isaaclab_rl/changelog.d/mhaiderbhai-zero-agent-actions.rst b/source/isaaclab_rl/changelog.d/mhaiderbhai-zero-agent-actions.rst new file mode 100644 index 000000000000..bfae95cb6773 --- /dev/null +++ b/source/isaaclab_rl/changelog.d/mhaiderbhai-zero-agent-actions.rst @@ -0,0 +1,5 @@ +Fixed +^^^^^ + +* Fixed the zero agent to use semantic neutral actions, support composite and multi-agent action spaces, + and reject invalid task configurations before launching the simulator. diff --git a/source/isaaclab_rl/isaaclab_rl/entrypoints/simple_agents.py b/source/isaaclab_rl/isaaclab_rl/entrypoints/simple_agents.py index fea6da73474f..b1b90c4717f7 100644 --- a/source/isaaclab_rl/isaaclab_rl/entrypoints/simple_agents.py +++ b/source/isaaclab_rl/isaaclab_rl/entrypoints/simple_agents.py @@ -6,7 +6,7 @@ """Checkpoint-free playback workflows for Isaac Lab environments. The zero and random agents are variations of playback that need no trained checkpoint: -the policy either emits constant zero actions or samples uniform random actions. +the policy either emits neutral actions or samples uniform random actions. """ from __future__ import annotations @@ -20,6 +20,7 @@ import torch from isaaclab.app import add_launcher_args, launch_simulation +from isaaclab.envs.utils.spaces import sample_space import isaaclab_tasks # noqa: F401 from isaaclab_tasks.utils import ( @@ -46,7 +47,7 @@ def run(argv: list[str] | None = None, *, policy: PolicyName) -> None: Args: argv: Command-line arguments excluding the executable name. Reads ``sys.argv`` when omitted. - policy: Action policy to apply, either constant zero actions or uniform random actions. + policy: Action policy to apply, either neutral actions or uniform random actions. Raises: ValueError: If the requested policy is not supported. @@ -61,13 +62,18 @@ def run(argv: list[str] | None = None, *, policy: PolicyName) -> None: # parse configuration via Hydra (supports preset selection, e.g. env.sim.physics=newton_mjwarp) env_cfg, _ = resolve_task_config(args_cli.task, "") - with launch_simulation(env_cfg, args_cli): - # override with CLI arguments - env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else env_cfg.scene.num_envs - env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device - if args_cli.disable_fabric: - env_cfg.sim.use_fabric = False + # override with CLI arguments and reject unsupported configurations before + # launching Kit or initializing a native physics backend. + env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else env_cfg.scene.num_envs + env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device + if args_cli.disable_fabric: + env_cfg.sim.use_fabric = False + try: + env_cfg.validate() + except (TypeError, ValueError) as exc: + raise SystemExit(f"Invalid environment configuration: {exc}") from None + with launch_simulation(env_cfg, args_cli): # create environment env = gym.make(args_cli.task, cfg=env_cfg) @@ -80,7 +86,6 @@ def run(argv: list[str] | None = None, *, policy: PolicyName) -> None: # keep running while any visualizer is open, and until the step budget is exhausted sim = env.unwrapped.sim device = env.unwrapped.device - zero_actions = torch.zeros(env.action_space.shape, device=device) step = 0 while sim.is_headless_or_exist_active_visualizer(): if args_cli.max_steps is not None and step >= args_cli.max_steps: @@ -89,7 +94,7 @@ def run(argv: list[str] | None = None, *, policy: PolicyName) -> None: # run everything in inference mode with torch.inference_mode(): if policy == "zero": - actions = zero_actions + actions = _get_neutral_actions(env) else: # sample actions from -1 to 1 actions = 2 * torch.rand(env.action_space.shape, device=device) - 1 @@ -99,6 +104,28 @@ def run(argv: list[str] | None = None, *, policy: PolicyName) -> None: env.close() +def _get_neutral_actions(env: gym.Env): + """Create semantically neutral actions for passive environment playback. + + Manager-based environments can provide semantic neutral actions for terms + where literal zeros are unsafe, such as absolute-pose IK. Direct-workflow + environments fall back to zero-filled samples of their declared Gymnasium + spaces, including composite and multi-agent spaces. + """ + unwrapped = env.unwrapped + action_manager = getattr(unwrapped, "action_manager", None) + if action_manager is not None: + return action_manager.neutral_actions + + if hasattr(unwrapped, "action_spaces"): + return { + agent: sample_space(space, unwrapped.device, batch_size=unwrapped.num_envs, fill_value=0) + for agent, space in unwrapped.action_spaces.items() + } + + return sample_space(unwrapped.single_action_space, unwrapped.device, batch_size=unwrapped.num_envs, fill_value=0) + + def _parse_args(argv: list[str] | None, policy: PolicyName) -> argparse.Namespace: """Parse the command line of a checkpoint-free agent and hand the remainder to Hydra.""" parser = argparse.ArgumentParser(description=_DESCRIPTIONS[policy]) diff --git a/source/isaaclab_rl/test/test_entrypoints.py b/source/isaaclab_rl/test/test_entrypoints.py index 39826af78c4a..a0ceb98525a8 100644 --- a/source/isaaclab_rl/test/test_entrypoints.py +++ b/source/isaaclab_rl/test/test_entrypoints.py @@ -11,11 +11,86 @@ import runpy import sys import types +from types import SimpleNamespace import gymnasium as gym +import numpy as np import pytest +import torch from isaaclab_rl.entrypoints import PlaybackRequest, TrainingRequest, api, dispatch +from isaaclab_rl.entrypoints import simple_agents as _simple_agents +from isaaclab_rl.entrypoints.simple_agents import _get_neutral_actions + + +def test_zero_agent_uses_manager_semantic_neutral_actions() -> None: + """The zero agent honors action-term neutral commands instead of forcing literal zeros.""" + expected = torch.tensor([[0.1, 0.2, 0.3, 1.0]]) + unwrapped = SimpleNamespace(action_manager=SimpleNamespace(neutral_actions=expected)) + + assert _get_neutral_actions(SimpleNamespace(unwrapped=unwrapped)) is expected + + +def test_zero_agent_supports_composite_direct_action_spaces() -> None: + """Direct environments receive tensorized zeros matching composite action spaces.""" + action_space = gym.spaces.Dict( + { + "continuous": gym.spaces.Box(low=-1.0, high=1.0, shape=(2,), dtype=np.float32), + "discrete": gym.spaces.Discrete(3), + } + ) + unwrapped = SimpleNamespace( + action_manager=None, + single_action_space=action_space, + device="cpu", + num_envs=2, + ) + + actions = _get_neutral_actions(SimpleNamespace(unwrapped=unwrapped)) + + assert torch.equal(actions["continuous"], torch.zeros(2, 2)) + assert torch.equal(actions["discrete"], torch.zeros(2, 1, dtype=torch.int64)) + + +def test_zero_agent_supports_direct_multi_agent_action_spaces() -> None: + """Direct multi-agent environments receive a zero action for every agent.""" + unwrapped = SimpleNamespace( + action_manager=None, + action_spaces={ + "robot": gym.spaces.Box(low=-1.0, high=1.0, shape=(2,), dtype=np.float32), + "object": gym.spaces.Discrete(2), + }, + device="cpu", + num_envs=3, + ) + + actions = _get_neutral_actions(SimpleNamespace(unwrapped=unwrapped)) + + assert torch.equal(actions["robot"], torch.zeros(3, 2)) + assert torch.equal(actions["object"], torch.zeros(3, 1, dtype=torch.int64)) + + +def test_zero_agent_rejects_invalid_config_before_launch(monkeypatch: pytest.MonkeyPatch) -> None: + """Unsupported task presets fail cleanly before a simulator backend is initialized.""" + + class _InvalidCfg: + scene = SimpleNamespace(num_envs=1) + sim = SimpleNamespace(device="cpu", use_fabric=True) + + def validate(self) -> None: + raise ValueError("unsupported physics backend") + + args = SimpleNamespace(num_envs=None, device=None, disable_fabric=False, task="Invalid-Task") + monkeypatch.setattr(_simple_agents, "_parse_args", lambda argv, policy: args) + monkeypatch.setattr(_simple_agents, "resolve_task_config", lambda task, agent: (_InvalidCfg(), None)) + monkeypatch.setattr( + _simple_agents, + "launch_simulation", + lambda *args, **kwargs: pytest.fail("simulation launched before config validation"), + ) + + with pytest.raises(SystemExit, match="Invalid environment configuration: unsupported physics backend"): + _simple_agents.run([], policy="zero") def test_train_request_adapts_typed_parameters_to_cli(monkeypatch) -> None: From 780dc7ed1120d9068ae02137dfdffdbffb08719f Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Mon, 31 Aug 2026 14:43:53 -0700 Subject: [PATCH 2/4] Let action terms handle zero actions --- .../tutorials/03_envs/create_cube_base_env.py | 4 +- .../mhaiderbhai-neutral-actions.minor.rst | 4 +- .../isaaclab/envs/manager_based_env.py | 7 +- .../isaaclab/envs/manager_based_rl_env.py | 7 +- .../envs/mdp/actions/binary_joint_actions.py | 8 +- .../envs/mdp/actions/joint_actions.py | 4 +- .../mdp/actions/joint_actions_to_limits.py | 6 +- .../envs/mdp/actions/non_holonomic_actions.py | 4 +- .../mdp/actions/pink_task_space_actions.py | 19 ++-- .../mdp/actions/rmpflow_task_space_actions.py | 19 ++-- .../mdp/actions/surface_gripper_actions.py | 4 +- .../envs/mdp/actions/task_space_actions.py | 86 ++++++++----------- .../isaaclab/managers/action_manager.py | 46 ++++------ .../check_manager_based_env_floating_cube.py | 4 +- .../test/envs/test_neutral_actions.py | 47 ++++++++-- .../test/envs/test_scale_randomization.py | 4 +- .../mhaiderbhai-zero-actions.minor.rst | 4 + .../mdp/actions/thrust_actions.py | 8 +- .../mhaiderbhai-zero-actions.minor.rst | 4 + .../envs/mdp/actions/newton_ik_actions.py | 70 ++++++++------- .../mhaiderbhai-zero-agent-actions.rst | 4 +- .../isaaclab_rl/entrypoints/simple_agents.py | 18 ++-- source/isaaclab_rl/test/test_entrypoints.py | 15 ++-- .../mhaiderbhai-zero-actions.minor.rst | 4 + .../contrib/franka_pour/mdp/actions.py | 4 +- .../locomanip_pick_place/mdp/actions.py | 7 +- .../mdp/pre_trained_policy_action.py | 4 +- .../contrib/ur10_particle_push/mdp/actions.py | 4 +- 28 files changed, 234 insertions(+), 185 deletions(-) create mode 100644 source/isaaclab_contrib/changelog.d/mhaiderbhai-zero-actions.minor.rst create mode 100644 source/isaaclab_newton/changelog.d/mhaiderbhai-zero-actions.minor.rst create mode 100644 source/isaaclab_tasks/changelog.d/mhaiderbhai-zero-actions.minor.rst diff --git a/scripts/tutorials/03_envs/create_cube_base_env.py b/scripts/tutorials/03_envs/create_cube_base_env.py index 43e269672344..68ae34edb93a 100644 --- a/scripts/tutorials/03_envs/create_cube_base_env.py +++ b/scripts/tutorials/03_envs/create_cube_base_env.py @@ -121,7 +121,9 @@ def processed_actions(self) -> torch.Tensor: Operations """ - def process_actions(self, actions: torch.Tensor): + def process_actions(self, actions: torch.Tensor | None): + if actions is None: + actions = self._asset.data.root_pos_w.torch - self._env.scene.env_origins # store the raw actions self._raw_actions[:] = actions # no-processing of actions diff --git a/source/isaaclab/changelog.d/mhaiderbhai-neutral-actions.minor.rst b/source/isaaclab/changelog.d/mhaiderbhai-neutral-actions.minor.rst index 9343ebf809c6..1bbf30477457 100644 --- a/source/isaaclab/changelog.d/mhaiderbhai-neutral-actions.minor.rst +++ b/source/isaaclab/changelog.d/mhaiderbhai-neutral-actions.minor.rst @@ -1,5 +1,5 @@ Added ^^^^^ -* Added semantic neutral actions for manager-based action terms, including absolute differential IK, - Pink IK, RMPFlow, and operational-space controllers. +* Allowed manager-based environments to accept ``None`` actions, with action terms applying their zero-action behavior. + Absolute differential IK, Pink IK, RMPFlow, and operational-space controllers hold their current poses. diff --git a/source/isaaclab/isaaclab/envs/manager_based_env.py b/source/isaaclab/isaaclab/envs/manager_based_env.py index e427b7d374be..38f2c3babde4 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_env.py +++ b/source/isaaclab/isaaclab/envs/manager_based_env.py @@ -500,7 +500,7 @@ def reset_to( # return observations return self.obs_buf, self.extras - def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]: + def step(self, action: torch.Tensor | None) -> tuple[VecEnvObs, dict]: """Execute one time-step of the environment's dynamics. The environment steps forward at a fixed time-step, while the physics simulation is @@ -521,13 +521,14 @@ def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]: app loop. Args: - action: The actions to apply on the environment. Shape is (num_envs, action_dim). + action: The actions to apply on the environment. Shape is (num_envs, action_dim). If None, each action + term applies its zero-action behavior. Returns: A tuple containing the observations and extras. """ # process actions - self.action_manager.process_action(action.to(self.device)) + self.action_manager.process_action(action.to(self.device) if action is not None else None) self.recorder_manager.record_pre_step() diff --git a/source/isaaclab/isaaclab/envs/manager_based_rl_env.py b/source/isaaclab/isaaclab/envs/manager_based_rl_env.py index ed751f1dd95e..00e858e8ff81 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_rl_env.py +++ b/source/isaaclab/isaaclab/envs/manager_based_rl_env.py @@ -173,7 +173,7 @@ def setup_manager_visualizers(self): Operations - MDP """ - def step(self, action: torch.Tensor) -> VecEnvStepReturn: + def step(self, action: torch.Tensor | None) -> VecEnvStepReturn: """Execute one time-step of the environment's dynamics and reset terminated environments. Unlike the :class:`ManagerBasedEnv.step` class, the function performs the following operations: @@ -199,13 +199,14 @@ def step(self, action: torch.Tensor) -> VecEnvStepReturn: - Post-reset re-renders for RTX sensors are also skipped. Args: - action: The actions to apply on the environment. Shape is (num_envs, action_dim). + action: The actions to apply on the environment. Shape is (num_envs, action_dim). If None, each action + term applies its zero-action behavior. Returns: A tuple containing the observations, rewards, resets (terminated and truncated) and extras. """ # process actions - self.action_manager.process_action(action.to(self.device)) + self.action_manager.process_action(action.to(self.device) if action is not None else None) self.recorder_manager.record_pre_step() diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/binary_joint_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/binary_joint_actions.py index db80bf2b6e64..9b49ec77317e 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/binary_joint_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/binary_joint_actions.py @@ -129,7 +129,9 @@ def IO_descriptor(self) -> GenericActionIODescriptor: Operations. """ - def process_actions(self, actions: torch.Tensor): + def process_actions(self, actions: torch.Tensor | None): + if actions is None: + actions = self._raw_actions.zero_() # store the raw actions self._raw_actions[:] = actions # compute the binary mask @@ -182,7 +184,9 @@ class AbsBinaryJointPositionAction(BinaryJointAction): cfg: actions_cfg.AbsBinaryJointPositionActionCfg """The configuration of the action term.""" - def process_actions(self, actions: torch.Tensor): + def process_actions(self, actions: torch.Tensor | None): + if actions is None: + actions = self._raw_actions.zero_() # store the raw actions self._raw_actions[:] = actions # compute the binary mask diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py index 435a04574a06..2d5a8cca4fe9 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py @@ -167,7 +167,9 @@ def IO_descriptor(self) -> GenericActionIODescriptor: Operations. """ - def process_actions(self, actions: torch.Tensor): + def process_actions(self, actions: torch.Tensor | None): + if actions is None: + actions = self._raw_actions.zero_() # store the raw actions self._raw_actions[:] = actions # apply the affine transformations diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions_to_limits.py b/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions_to_limits.py index f6c416b97337..e6ba2867f987 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions_to_limits.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions_to_limits.py @@ -152,7 +152,9 @@ def IO_descriptor(self) -> GenericActionIODescriptor: Operations. """ - def process_actions(self, actions: torch.Tensor): + def process_actions(self, actions: torch.Tensor | None): + if actions is None: + actions = self._raw_actions.zero_() # store the raw actions self._raw_actions[:] = actions # apply affine transformations @@ -279,7 +281,7 @@ def reset(self, env_ids: Sequence[int] | None = None) -> None: ) self._prev_applied_actions[env_ids, :] = curr_applied_actions - def process_actions(self, actions: torch.Tensor): + def process_actions(self, actions: torch.Tensor | None): # apply affine transformations super().process_actions(actions) # set position targets as moving average diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/non_holonomic_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/non_holonomic_actions.py index 08d5e9be1c85..7847493a5136 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/non_holonomic_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/non_holonomic_actions.py @@ -172,7 +172,9 @@ def IO_descriptor(self) -> GenericActionIODescriptor: Operations. """ - def process_actions(self, actions): + def process_actions(self, actions: torch.Tensor | None): + if actions is None: + actions = self._raw_actions.zero_() # store the raw actions self._raw_actions[:] = actions self._processed_actions = self.raw_actions * self._scale + self._offset diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py index 32949738f918..2590361d997b 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py @@ -170,14 +170,6 @@ def processed_actions(self) -> torch.Tensor: """Get the processed actions tensor.""" return self._processed_actions - @property - def neutral_actions(self) -> torch.Tensor: - """Raw actions that hold the controlled frames and hand joints at their current state.""" - frame_poses = self._asset.data.body_link_pose_w.torch[:, self._controlled_frame_ids].clone() - frame_poses[..., :3] -= self._env.scene.env_origins.unsqueeze(1) - hand_joint_positions = self._asset.data.joint_pos.torch[:, self._hand_joint_ids] - return torch.cat((frame_poses.flatten(start_dim=1), hand_joint_positions), dim=-1) - @property def IO_descriptor(self) -> GenericActionIODescriptor: """The IO descriptor of the action term. @@ -207,12 +199,19 @@ def IO_descriptor(self) -> GenericActionIODescriptor: # Operations. # """ - def process_actions(self, actions: torch.Tensor) -> None: + def process_actions(self, actions: torch.Tensor | None) -> None: """Process the input actions and set targets for each task. Args: - actions: The input actions tensor. + actions: The input actions tensor. If None, the current controlled-frame poses and hand-joint positions + are used. """ + if actions is None: + frame_poses = self._asset.data.body_link_pose_w.torch[:, self._controlled_frame_ids].clone() + frame_poses[..., :3] -= self._env.scene.env_origins.unsqueeze(1) + hand_joint_positions = self._asset.data.joint_pos.torch[:, self._hand_joint_ids] + actions = torch.cat((frame_poses.flatten(start_dim=1), hand_joint_positions), dim=-1) + # Store raw actions self._raw_actions[:] = actions diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/rmpflow_task_space_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/rmpflow_task_space_actions.py index 1893cc465f78..93a7fb09bbce 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/rmpflow_task_space_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/rmpflow_task_space_actions.py @@ -118,16 +118,6 @@ def raw_actions(self) -> torch.Tensor: def processed_actions(self) -> torch.Tensor: return self._processed_actions - @property - def neutral_actions(self) -> torch.Tensor: - """Raw actions that hold the current end-effector pose.""" - if self.cfg.use_relative_mode: - return super().neutral_actions - - ee_pos, ee_quat = self._compute_frame_pose() - command = torch.cat((ee_pos, ee_quat), dim=-1) - return torch.where(self._scale != 0.0, command / self._scale, torch.zeros_like(command)) - @property def jacobian_w(self) -> torch.Tensor: return self._asset.data.body_link_jacobian_w.torch[:, self._jacobi_body_idx, :, self._jacobi_joint_ids] @@ -146,7 +136,14 @@ def jacobian_b(self) -> torch.Tensor: """ # This is called each env.step() - def process_actions(self, actions: torch.Tensor): + def process_actions(self, actions: torch.Tensor | None): + if actions is None: + if self.cfg.use_relative_mode: + actions = self._raw_actions.zero_() + else: + ee_pos, ee_quat = self._compute_frame_pose() + command = torch.cat((ee_pos, ee_quat), dim=-1) + actions = torch.where(self._scale != 0.0, command / self._scale, torch.zeros_like(command)) # store the raw actions self._raw_actions[:] = actions self._processed_actions[:] = self.raw_actions * self._scale diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/surface_gripper_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/surface_gripper_actions.py index 699743eb918b..62d4ff566cff 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/surface_gripper_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/surface_gripper_actions.py @@ -85,7 +85,9 @@ def processed_actions(self) -> torch.Tensor: Operations. """ - def process_actions(self, actions: torch.Tensor): + def process_actions(self, actions: torch.Tensor | None): + if actions is None: + actions = self._raw_actions.zero_() # store the raw actions self._raw_actions[:] = actions # compute the binary mask diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py index 33adfba89b32..a9e9f5c000dc 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py @@ -139,16 +139,6 @@ def raw_actions(self) -> torch.Tensor: def processed_actions(self) -> torch.Tensor: return self._processed_actions - @property - def neutral_actions(self) -> torch.Tensor: - """Raw actions that hold the current end-effector pose.""" - if self.cfg.controller.use_relative_mode: - return super().neutral_actions - - ee_pos, ee_quat = self._compute_frame_pose() - command = ee_pos if self.cfg.controller.command_type == "position" else torch.cat((ee_pos, ee_quat), dim=-1) - return torch.where(self._scale != 0.0, command / self._scale, torch.zeros_like(command)) - @property def jacobian_w(self) -> torch.Tensor: return self._asset.data.body_link_jacobian_w.torch[:, self._jacobi_body_idx, :, self._jacobi_joint_ids] @@ -197,7 +187,16 @@ def IO_descriptor(self) -> GenericActionIODescriptor: Operations. """ - def process_actions(self, actions: torch.Tensor): + def process_actions(self, actions: torch.Tensor | None): + if actions is None: + if self.cfg.controller.use_relative_mode: + actions = self._raw_actions.zero_() + else: + ee_pos, ee_quat = self._compute_frame_pose() + command = ( + ee_pos if self.cfg.controller.command_type == "position" else torch.cat((ee_pos, ee_quat), dim=-1) + ) + actions = torch.where(self._scale != 0.0, command / self._scale, torch.zeros_like(command)) # store the raw actions self._raw_actions[:] = actions self._processed_actions[:] = self.raw_actions * self._scale @@ -460,40 +459,6 @@ def processed_actions(self) -> torch.Tensor: """Processed actions for operational space control.""" return self._processed_actions - @property - def neutral_actions(self) -> torch.Tensor: - """Raw actions that hold the current end-effector pose and apply no wrench.""" - actions = super().neutral_actions - if self._pose_abs_idx is None: - return actions - - self._compute_ee_pose() - self._compute_task_frame_pose() - if self._task_frame_pose_b is None: - ee_pos_task = self._ee_pose_b[:, :3] - ee_quat_task = self._ee_pose_b[:, 3:7] - else: - ee_pos_task, ee_quat_task = math_utils.subtract_frame_transforms( - self._task_frame_pose_b[:, :3], - self._task_frame_pose_b[:, 3:7], - self._ee_pose_b[:, :3], - self._ee_pose_b[:, 3:7], - ) - - position_slice = slice(self._pose_abs_idx, self._pose_abs_idx + 3) - orientation_slice = slice(self._pose_abs_idx + 3, self._pose_abs_idx + 7) - actions[:, position_slice] = torch.where( - self._position_scale != 0.0, - ee_pos_task / self._position_scale, - torch.zeros_like(ee_pos_task), - ) - actions[:, orientation_slice] = torch.where( - self._orientation_scale != 0.0, - ee_quat_task / self._orientation_scale, - torch.zeros_like(ee_quat_task), - ) - return actions - @property def jacobian_w(self) -> torch.Tensor: return self._asset.data.body_link_jacobian_w.torch[:, self._jacobi_ee_body_idx, :, self._jacobi_joint_idx] @@ -552,12 +517,13 @@ def IO_descriptor(self) -> GenericActionIODescriptor: Operations. """ - def process_actions(self, actions: torch.Tensor): + def process_actions(self, actions: torch.Tensor | None): """Pre-processes the raw actions and sets them as commands for for operational space control. Args: - actions (torch.Tensor): The raw actions for operational space control. It is a tensor of - shape (``num_envs``, ``action_dim``). + actions: The raw actions for operational space control. It is a tensor of shape + (``num_envs``, ``action_dim``). If None, the controller holds the current absolute pose and applies + zero relative pose and wrench commands. """ # Update ee pose, which would be used by relative targets (i.e., pose_rel) @@ -566,6 +532,30 @@ def process_actions(self, actions: torch.Tensor): # Update task frame pose w.r.t. the root frame. self._compute_task_frame_pose() + if actions is None: + actions = self._raw_actions.zero_() + if self._pose_abs_idx is not None: + if self._task_frame_pose_b is None: + ee_pos_task = self._ee_pose_b[:, :3] + ee_quat_task = self._ee_pose_b[:, 3:7] + else: + ee_pos_task, ee_quat_task = math_utils.subtract_frame_transforms( + self._task_frame_pose_b[:, :3], + self._task_frame_pose_b[:, 3:7], + self._ee_pose_b[:, :3], + self._ee_pose_b[:, 3:7], + ) + actions[:, self._pose_abs_idx : self._pose_abs_idx + 3] = torch.where( + self._position_scale != 0.0, + ee_pos_task / self._position_scale, + torch.zeros_like(ee_pos_task), + ) + actions[:, self._pose_abs_idx + 3 : self._pose_abs_idx + 7] = torch.where( + self._orientation_scale != 0.0, + ee_quat_task / self._orientation_scale, + torch.zeros_like(ee_quat_task), + ) + # Pre-process the raw actions for operational space control. self._preprocess_actions(actions) diff --git a/source/isaaclab/isaaclab/managers/action_manager.py b/source/isaaclab/isaaclab/managers/action_manager.py index 1e952628b168..c58b5df70ba7 100644 --- a/source/isaaclab/isaaclab/managers/action_manager.py +++ b/source/isaaclab/isaaclab/managers/action_manager.py @@ -88,15 +88,6 @@ def processed_actions(self) -> torch.Tensor: """The actions computed by the term after applying any processing.""" raise NotImplementedError - @property - def neutral_actions(self) -> torch.Tensor: - """Raw actions suitable for passive agent playback. - - The default is a zero-filled tensor. Action terms for which zero has a different or invalid meaning, - such as absolute-pose controllers, should override this property with a semantically neutral command. - """ - return torch.zeros_like(self.raw_actions) - @property def has_debug_vis_implementation(self) -> bool: """Whether the action term has a debug visualization implemented.""" @@ -148,14 +139,14 @@ def set_debug_vis(self, debug_vis: bool) -> bool: return True @abstractmethod - def process_actions(self, actions: torch.Tensor): + def process_actions(self, actions: torch.Tensor | None): """Processes the actions sent to the environment. Note: This function is called once per environment step by the manager. Args: - actions: The actions to process. + actions: The actions to process. If None, the action term applies its zero-action behavior. """ raise NotImplementedError @@ -272,18 +263,6 @@ def prev_action(self) -> torch.Tensor: """The previous actions sent to the environment. Shape is (num_envs, total_action_dim).""" return self._prev_action - @property - def neutral_actions(self) -> torch.Tensor: - """Raw actions suitable for passive playback of all active action terms. - - The returned tensor has shape ``(num_envs, total_action_dim)``. Since - some terms derive their neutral command from the current simulation - state, consumers should retrieve this property immediately before use. - """ - if not self._terms: - return torch.zeros_like(self._action) - return torch.cat([term.neutral_actions for term in self._terms.values()], dim=-1) - @property def has_debug_vis_implementation(self) -> bool: """Whether the command terms have debug visualization implemented.""" @@ -384,26 +363,31 @@ def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, torch.Tensor] # nothing to log here return {} - def process_action(self, action: torch.Tensor): + def process_action(self, action: torch.Tensor | None): """Processes the actions sent to the environment. Note: This function should be called once per environment step. Args: - action: The actions to process. + action: The actions to process. If None, each action term applies its zero-action behavior. """ - # check if action dimension is valid - if self.total_action_dim != action.shape[1]: - raise ValueError(f"Invalid action shape, expected: {self.total_action_dim}, received: {action.shape[1]}.") - # store the input actions self._prev_action[:] = self._action - self._action[:] = action.to(self.device) + if action is None: + self._action.zero_() + else: + # check if action dimension is valid + if self.total_action_dim != action.shape[1]: + raise ValueError( + f"Invalid action shape, expected: {self.total_action_dim}, received: {action.shape[1]}." + ) + # store the input actions + self._action[:] = action.to(self.device) # split the actions and apply to each tensor idx = 0 for term in self._terms.values(): - term_actions = action[:, idx : idx + term.action_dim] + term_actions = None if action is None else self._action[:, idx : idx + term.action_dim] term.process_actions(term_actions) idx += term.action_dim diff --git a/source/isaaclab/test/envs/check_manager_based_env_floating_cube.py b/source/isaaclab/test/envs/check_manager_based_env_floating_cube.py index ef0a151434b5..efb22d327dd9 100644 --- a/source/isaaclab/test/envs/check_manager_based_env_floating_cube.py +++ b/source/isaaclab/test/envs/check_manager_based_env_floating_cube.py @@ -120,7 +120,9 @@ def processed_actions(self) -> torch.Tensor: Operations """ - def process_actions(self, actions: torch.Tensor): + def process_actions(self, actions: torch.Tensor | None): + if actions is None: + actions = self._asset.data.root_pos_w.torch - self._env.scene.env_origins # store the raw actions self._raw_actions[:] = actions # no-processing of actions diff --git a/source/isaaclab/test/envs/test_neutral_actions.py b/source/isaaclab/test/envs/test_neutral_actions.py index d0b6708c314d..2af8b5296b75 100644 --- a/source/isaaclab/test/envs/test_neutral_actions.py +++ b/source/isaaclab/test/envs/test_neutral_actions.py @@ -3,17 +3,44 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Tests for semantic neutral actions.""" +"""Tests for action-term zero actions.""" from types import SimpleNamespace import torch from isaaclab.envs.mdp.actions.pink_task_space_actions import PinkInverseKinematicsAction +from isaaclab.managers.action_manager import ActionManager -def test_pink_neutral_actions_use_current_frame_poses() -> None: - """Pink IK neutral actions contain valid current poses and hand joint positions.""" +def test_action_manager_dispatches_none_and_records_zero_action() -> None: + """The action manager lets terms resolve None while recording a conceptual zero action.""" + + class _ActionTerm: + action_dim = 2 + raw_actions = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + received_none = False + + def process_actions(self, actions: torch.Tensor | None) -> None: + assert actions is None + self.received_none = True + + manager = object.__new__(ActionManager) + term = _ActionTerm() + manager._terms = {"term": term} + manager._action = torch.full((2, 2), 5.0) + manager._prev_action = torch.full((2, 2), -1.0) + manager._resolve_terms_handle = None + + manager.process_action(None) + + assert term.received_none + assert torch.equal(manager.action, torch.zeros(2, 2)) + assert torch.equal(manager.prev_action, torch.full((2, 2), 5.0)) + + +def test_pink_none_actions_use_current_frame_poses() -> None: + """Pink IK resolves None to valid current poses and hand joint positions.""" action_term = object.__new__(PinkInverseKinematicsAction) action_term._controlled_frame_ids = [1, 0] action_term._hand_joint_ids = [1, 3] @@ -33,11 +60,19 @@ def test_pink_neutral_actions_use_current_frame_poses() -> None: ) ) action_term._env = SimpleNamespace(scene=SimpleNamespace(env_origins=env_origins)) + action_term.cfg = SimpleNamespace(controller=SimpleNamespace(num_hand_joints=2)) + action_term._raw_actions = torch.zeros(2, 16) + action_term._get_base_link_frame_transform = lambda: torch.eye(4).repeat(2, 1, 1) + action_term._extract_controlled_frame_poses = lambda actions: actions[:, :14] + action_term._transform_poses_to_base_link_frame = lambda poses: poses + action_term._set_task_targets = lambda poses: None - actions = action_term.neutral_actions + action_term.process_actions(None) expected_poses = body_poses[:, [1, 0]].clone() expected_poses[..., :3] -= env_origins.unsqueeze(1) expected = torch.cat((expected_poses.flatten(start_dim=1), joint_positions[:, [1, 3]]), dim=-1) - assert torch.equal(actions, expected) - assert torch.all(torch.linalg.vector_norm(actions[:, :14].reshape(2, 2, 7)[..., 3:7], dim=-1) == 1.0) + assert torch.equal(action_term.raw_actions, expected) + assert torch.all( + torch.linalg.vector_norm(action_term.raw_actions[:, :14].reshape(2, 2, 7)[..., 3:7], dim=-1) == 1.0 + ) diff --git a/source/isaaclab/test/envs/test_scale_randomization.py b/source/isaaclab/test/envs/test_scale_randomization.py index d88029a8965e..c007e754eb1f 100644 --- a/source/isaaclab/test/envs/test_scale_randomization.py +++ b/source/isaaclab/test/envs/test_scale_randomization.py @@ -95,7 +95,9 @@ def processed_actions(self) -> torch.Tensor: Operations """ - def process_actions(self, actions: torch.Tensor): + def process_actions(self, actions: torch.Tensor | None): + if actions is None: + actions = self._asset.data.root_pos_w.torch - self._env.scene.env_origins # store the raw actions self._raw_actions[:] = actions # no-processing of actions diff --git a/source/isaaclab_contrib/changelog.d/mhaiderbhai-zero-actions.minor.rst b/source/isaaclab_contrib/changelog.d/mhaiderbhai-zero-actions.minor.rst new file mode 100644 index 000000000000..715e429a6a9d --- /dev/null +++ b/source/isaaclab_contrib/changelog.d/mhaiderbhai-zero-actions.minor.rst @@ -0,0 +1,4 @@ +Added +^^^^^ + +* Added support for ``None`` zero actions to multirotor thrust and navigation action terms. diff --git a/source/isaaclab_contrib/isaaclab_contrib/mdp/actions/thrust_actions.py b/source/isaaclab_contrib/isaaclab_contrib/mdp/actions/thrust_actions.py index 897d621246b9..852b765641fe 100644 --- a/source/isaaclab_contrib/isaaclab_contrib/mdp/actions/thrust_actions.py +++ b/source/isaaclab_contrib/isaaclab_contrib/mdp/actions/thrust_actions.py @@ -200,7 +200,7 @@ def reset(self, env_ids: Sequence[int] | None = None) -> None: """ self._raw_actions[env_ids] = 0.0 - def process_actions(self, actions: torch.Tensor): + def process_actions(self, actions: torch.Tensor | None): r"""Process actions by applying scaling, offset, and clipping. This method transforms raw policy actions into thrust commands through @@ -216,12 +216,14 @@ def process_actions(self, actions: torch.Tensor): Args: actions: Raw action tensor from the policy. Shape is ``(num_envs, action_dim)``. - Typically in the range [-1, 1] for normalized policies. + Typically in the range [-1, 1] for normalized policies. If None, zeros are used. Note: The processed actions are stored internally and applied during the next :meth:`apply_actions` call. """ + if actions is None: + actions = self._raw_actions.zero_() # store the raw actions self._raw_actions[:] = actions # apply the affine transformations @@ -331,7 +333,7 @@ def IO_descriptor(self) -> GenericActionIODescriptor: descriptor.action_type = "NavigationAction" return descriptor - def process_actions(self, actions: torch.Tensor): + def process_actions(self, actions: torch.Tensor | None): """Process actions by applying scaling, offset, and clipping.""" # Call parent to handle basic processing super().process_actions(actions) diff --git a/source/isaaclab_newton/changelog.d/mhaiderbhai-zero-actions.minor.rst b/source/isaaclab_newton/changelog.d/mhaiderbhai-zero-actions.minor.rst new file mode 100644 index 000000000000..a548ecbd5554 --- /dev/null +++ b/source/isaaclab_newton/changelog.d/mhaiderbhai-zero-actions.minor.rst @@ -0,0 +1,4 @@ +Added +^^^^^ + +* Added support for ``None`` zero actions to Newton inverse-kinematics action terms, which hold their current poses. diff --git a/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py b/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py index 498aef87ab19..5fc3aed004d2 100644 --- a/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py +++ b/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py @@ -50,6 +50,7 @@ def _ik_world_target_kernel( scale: wp.array(dtype=wp.float32), command_code: int, use_relative: int, + zero_action: int, out_pos: wp.array(dtype=wp.vec3f), out_rot: wp.array(dtype=wp.vec4f), ): @@ -71,42 +72,43 @@ def _ik_world_target_kernel( target_pos = ee_pos target_rot = ee_rot - if command_code == 0: # COMMAND_POSITION - disp = wp.vec3f( - action[i, action_offset + 0] * scale[0], - action[i, action_offset + 1] * scale[1], - action[i, action_offset + 2] * scale[2], - ) - target_pos = ee_pos + disp if use_relative == 1 else disp - else: - if use_relative == 1: - target_pos = ee_pos + wp.vec3f( + if zero_action == 0: + if command_code == 0: # COMMAND_POSITION + disp = wp.vec3f( action[i, action_offset + 0] * scale[0], action[i, action_offset + 1] * scale[1], action[i, action_offset + 2] * scale[2], ) - rot_vec = wp.vec3f( - action[i, action_offset + 3] * scale[3], - action[i, action_offset + 4] * scale[4], - action[i, action_offset + 5] * scale[5], - ) - angle = wp.length(rot_vec) - delta_rot = wp.quat_identity() - if angle > 1.0e-6: - delta_rot = wp.quat_from_axis_angle(rot_vec / angle, angle) - target_rot = delta_rot * ee_rot + target_pos = ee_pos + disp if use_relative == 1 else disp else: - target_pos = wp.vec3f( - action[i, action_offset + 0] * scale[0], - action[i, action_offset + 1] * scale[1], - action[i, action_offset + 2] * scale[2], - ) - target_rot = wp.quatf( - action[i, action_offset + 3] * scale[3], - action[i, action_offset + 4] * scale[4], - action[i, action_offset + 5] * scale[5], - action[i, action_offset + 6] * scale[6], - ) + if use_relative == 1: + target_pos = ee_pos + wp.vec3f( + action[i, action_offset + 0] * scale[0], + action[i, action_offset + 1] * scale[1], + action[i, action_offset + 2] * scale[2], + ) + rot_vec = wp.vec3f( + action[i, action_offset + 3] * scale[3], + action[i, action_offset + 4] * scale[4], + action[i, action_offset + 5] * scale[5], + ) + angle = wp.length(rot_vec) + delta_rot = wp.quat_identity() + if angle > 1.0e-6: + delta_rot = wp.quat_from_axis_angle(rot_vec / angle, angle) + target_rot = delta_rot * ee_rot + else: + target_pos = wp.vec3f( + action[i, action_offset + 0] * scale[0], + action[i, action_offset + 1] * scale[1], + action[i, action_offset + 2] * scale[2], + ) + target_rot = wp.quatf( + action[i, action_offset + 3] * scale[3], + action[i, action_offset + 4] * scale[4], + action[i, action_offset + 5] * scale[5], + action[i, action_offset + 6] * scale[6], + ) # Broadcast against the env-0 prototype root (all roots identical, validated). world_t = wp.transform_multiply(wp.transformf(root_pos_w[0], root_quat_w[0]), wp.transformf(target_pos, target_rot)) @@ -266,7 +268,10 @@ def IO_descriptor(self) -> GenericActionIODescriptor: self._IO_descriptor.extras["coordinate_names"] = self._action_coordinate_names() return self._IO_descriptor - def process_actions(self, actions: torch.Tensor) -> None: + def process_actions(self, actions: torch.Tensor | None) -> None: + zero_action = actions is None + if actions is None: + actions = self._raw_actions.zero_() self._raw_actions[:] = actions self._processed_actions[:] = self._raw_actions if self._clip is not None: @@ -298,6 +303,7 @@ def process_actions(self, actions: torch.Tensor) -> None: obj.scale, obj.command_code, obj.use_relative, + int(zero_action), obj.position_objective.target_positions, obj.rotation_objective.target_rotations, ], diff --git a/source/isaaclab_rl/changelog.d/mhaiderbhai-zero-agent-actions.rst b/source/isaaclab_rl/changelog.d/mhaiderbhai-zero-agent-actions.rst index bfae95cb6773..68fb064e2a90 100644 --- a/source/isaaclab_rl/changelog.d/mhaiderbhai-zero-agent-actions.rst +++ b/source/isaaclab_rl/changelog.d/mhaiderbhai-zero-agent-actions.rst @@ -1,5 +1,5 @@ Fixed ^^^^^ -* Fixed the zero agent to use semantic neutral actions, support composite and multi-agent action spaces, - and reject invalid task configurations before launching the simulator. +* Fixed the zero agent to defer zero-action handling to manager-based action terms, support composite and multi-agent + action spaces, and reject invalid task configurations before launching the simulator. diff --git a/source/isaaclab_rl/isaaclab_rl/entrypoints/simple_agents.py b/source/isaaclab_rl/isaaclab_rl/entrypoints/simple_agents.py index b1b90c4717f7..7e91ff5bf9b2 100644 --- a/source/isaaclab_rl/isaaclab_rl/entrypoints/simple_agents.py +++ b/source/isaaclab_rl/isaaclab_rl/entrypoints/simple_agents.py @@ -6,7 +6,7 @@ """Checkpoint-free playback workflows for Isaac Lab environments. The zero and random agents are variations of playback that need no trained checkpoint: -the policy either emits neutral actions or samples uniform random actions. +the policy either requests each action term's zero-action behavior or samples uniform random actions. """ from __future__ import annotations @@ -47,7 +47,7 @@ def run(argv: list[str] | None = None, *, policy: PolicyName) -> None: Args: argv: Command-line arguments excluding the executable name. Reads ``sys.argv`` when omitted. - policy: Action policy to apply, either neutral actions or uniform random actions. + policy: Action policy to apply, either term-specific zero actions or uniform random actions. Raises: ValueError: If the requested policy is not supported. @@ -94,7 +94,7 @@ def run(argv: list[str] | None = None, *, policy: PolicyName) -> None: # run everything in inference mode with torch.inference_mode(): if policy == "zero": - actions = _get_neutral_actions(env) + actions = _get_zero_actions(env) else: # sample actions from -1 to 1 actions = 2 * torch.rand(env.action_space.shape, device=device) - 1 @@ -104,18 +104,16 @@ def run(argv: list[str] | None = None, *, policy: PolicyName) -> None: env.close() -def _get_neutral_actions(env: gym.Env): - """Create semantically neutral actions for passive environment playback. +def _get_zero_actions(env: gym.Env): + """Create zero actions for passive environment playback. - Manager-based environments can provide semantic neutral actions for terms - where literal zeros are unsafe, such as absolute-pose IK. Direct-workflow - environments fall back to zero-filled samples of their declared Gymnasium - spaces, including composite and multi-agent spaces. + Manager-based environments accept None so that each action term can apply its zero-action behavior. Direct-workflow + environments use zero-filled samples of their declared Gymnasium spaces, including composite and multi-agent spaces. """ unwrapped = env.unwrapped action_manager = getattr(unwrapped, "action_manager", None) if action_manager is not None: - return action_manager.neutral_actions + return None if hasattr(unwrapped, "action_spaces"): return { diff --git a/source/isaaclab_rl/test/test_entrypoints.py b/source/isaaclab_rl/test/test_entrypoints.py index a0ceb98525a8..826f2f161986 100644 --- a/source/isaaclab_rl/test/test_entrypoints.py +++ b/source/isaaclab_rl/test/test_entrypoints.py @@ -20,15 +20,14 @@ from isaaclab_rl.entrypoints import PlaybackRequest, TrainingRequest, api, dispatch from isaaclab_rl.entrypoints import simple_agents as _simple_agents -from isaaclab_rl.entrypoints.simple_agents import _get_neutral_actions +from isaaclab_rl.entrypoints.simple_agents import _get_zero_actions -def test_zero_agent_uses_manager_semantic_neutral_actions() -> None: - """The zero agent honors action-term neutral commands instead of forcing literal zeros.""" - expected = torch.tensor([[0.1, 0.2, 0.3, 1.0]]) - unwrapped = SimpleNamespace(action_manager=SimpleNamespace(neutral_actions=expected)) +def test_zero_agent_defers_to_manager_action_terms() -> None: + """The zero agent lets manager-based action terms define their zero-action behavior.""" + unwrapped = SimpleNamespace(action_manager=SimpleNamespace()) - assert _get_neutral_actions(SimpleNamespace(unwrapped=unwrapped)) is expected + assert _get_zero_actions(SimpleNamespace(unwrapped=unwrapped)) is None def test_zero_agent_supports_composite_direct_action_spaces() -> None: @@ -46,7 +45,7 @@ def test_zero_agent_supports_composite_direct_action_spaces() -> None: num_envs=2, ) - actions = _get_neutral_actions(SimpleNamespace(unwrapped=unwrapped)) + actions = _get_zero_actions(SimpleNamespace(unwrapped=unwrapped)) assert torch.equal(actions["continuous"], torch.zeros(2, 2)) assert torch.equal(actions["discrete"], torch.zeros(2, 1, dtype=torch.int64)) @@ -64,7 +63,7 @@ def test_zero_agent_supports_direct_multi_agent_action_spaces() -> None: num_envs=3, ) - actions = _get_neutral_actions(SimpleNamespace(unwrapped=unwrapped)) + actions = _get_zero_actions(SimpleNamespace(unwrapped=unwrapped)) assert torch.equal(actions["robot"], torch.zeros(3, 2)) assert torch.equal(actions["object"], torch.zeros(3, 1, dtype=torch.int64)) diff --git a/source/isaaclab_tasks/changelog.d/mhaiderbhai-zero-actions.minor.rst b/source/isaaclab_tasks/changelog.d/mhaiderbhai-zero-actions.minor.rst new file mode 100644 index 000000000000..ee3dd361456d --- /dev/null +++ b/source/isaaclab_tasks/changelog.d/mhaiderbhai-zero-actions.minor.rst @@ -0,0 +1,4 @@ +Added +^^^^^ + +* Added support for ``None`` zero actions to task-specific action terms. diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/actions.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/actions.py index 4e2491b4832a..5d44a42ec763 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/actions.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/actions.py @@ -39,7 +39,7 @@ def __init__(self, cfg: EMARelativeJointPositionActionCfg, env: ManagerBasedEnv) raise ValueError(f"Moving-average weight must lie in (0, 1], got {self._alpha}.") self._previous_delta = torch.zeros_like(self._processed_actions) - def process_actions(self, actions: torch.Tensor) -> None: + def process_actions(self, actions: torch.Tensor | None) -> None: """Affine-map the raw action, then smooth only the commanded joint delta.""" super().process_actions(actions) self._processed_actions.lerp_(self._previous_delta, 1.0 - self._alpha) @@ -136,7 +136,7 @@ def set_reset_position( expanded = position.expand(-1, self._num_joints) self._processed_actions[selected] = expanded - def process_actions(self, actions: torch.Tensor) -> None: + def process_actions(self, actions: torch.Tensor | None) -> None: previous_target = self._processed_actions super().process_actions(actions) previous_target.lerp_(self._processed_actions, self._alpha) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/locomanip_pick_place/mdp/actions.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/locomanip_pick_place/mdp/actions.py index 87996ca9b066..70678632ec3d 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/locomanip_pick_place/mdp/actions.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/locomanip_pick_place/mdp/actions.py @@ -92,13 +92,16 @@ def _compose_policy_input(self, base_command: torch.Tensor, obs_tensor: torch.Te return policy_input - def process_actions(self, actions: torch.Tensor): + def process_actions(self, actions: torch.Tensor | None): """Process the input actions using the locomotion policy. Args: - actions: The lower body commands. + actions: The lower body commands. If None, zeros are used. """ + if actions is None: + actions = torch.zeros((self.num_envs, self.action_dim), device=self.device) + # Extract base command from the action tensor # Assuming the base command [vx, vy, wz, hip_height] base_command = actions diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/navigation/mdp/pre_trained_policy_action.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/navigation/mdp/pre_trained_policy_action.py index 4857d63711e1..61073a9dc002 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/navigation/mdp/pre_trained_policy_action.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/navigation/mdp/pre_trained_policy_action.py @@ -90,7 +90,9 @@ def processed_actions(self) -> torch.Tensor: Operations. """ - def process_actions(self, actions: torch.Tensor): + def process_actions(self, actions: torch.Tensor | None): + if actions is None: + actions = self._raw_actions.zero_() self._raw_actions[:] = actions def apply_actions(self): diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/ur10_particle_push/mdp/actions.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/ur10_particle_push/mdp/actions.py index a8f01c5d6220..7060053ae900 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/ur10_particle_push/mdp/actions.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/ur10_particle_push/mdp/actions.py @@ -50,8 +50,10 @@ def invalid_actions(self) -> torch.Tensor: """Whether the latest policy action contained a non-finite component.""" return self._invalid_actions - def process_actions(self, actions: torch.Tensor) -> None: + def process_actions(self, actions: torch.Tensor | None) -> None: """Sanitize the policy action and construct one bounded joint target.""" + if actions is None: + actions = torch.zeros_like(self._raw_actions) self._previous_actions.copy_(self._raw_actions) self._invalid_actions.copy_(~torch.isfinite(actions).all(dim=1)) self._raw_actions.copy_(torch.nan_to_num(actions, nan=0.0, posinf=1.0, neginf=-1.0).clamp(-1.0, 1.0)) From 0f1eb1a8839021ca76dc2ab338cf1a89a94eb2ce Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Mon, 31 Aug 2026 18:58:09 -0700 Subject: [PATCH 3/4] Revert "Let action terms handle zero actions" This reverts commit 780dc7ed1120d9068ae02137dfdffdbffb08719f. --- .../tutorials/03_envs/create_cube_base_env.py | 4 +- .../mhaiderbhai-neutral-actions.minor.rst | 4 +- .../isaaclab/envs/manager_based_env.py | 7 +- .../isaaclab/envs/manager_based_rl_env.py | 7 +- .../envs/mdp/actions/binary_joint_actions.py | 8 +- .../envs/mdp/actions/joint_actions.py | 4 +- .../mdp/actions/joint_actions_to_limits.py | 6 +- .../envs/mdp/actions/non_holonomic_actions.py | 4 +- .../mdp/actions/pink_task_space_actions.py | 19 ++-- .../mdp/actions/rmpflow_task_space_actions.py | 19 ++-- .../mdp/actions/surface_gripper_actions.py | 4 +- .../envs/mdp/actions/task_space_actions.py | 86 +++++++++++-------- .../isaaclab/managers/action_manager.py | 46 ++++++---- .../check_manager_based_env_floating_cube.py | 4 +- .../test/envs/test_neutral_actions.py | 47 ++-------- .../test/envs/test_scale_randomization.py | 4 +- .../mhaiderbhai-zero-actions.minor.rst | 4 - .../mdp/actions/thrust_actions.py | 8 +- .../mhaiderbhai-zero-actions.minor.rst | 4 - .../envs/mdp/actions/newton_ik_actions.py | 70 +++++++-------- .../mhaiderbhai-zero-agent-actions.rst | 4 +- .../isaaclab_rl/entrypoints/simple_agents.py | 18 ++-- source/isaaclab_rl/test/test_entrypoints.py | 15 ++-- .../mhaiderbhai-zero-actions.minor.rst | 4 - .../contrib/franka_pour/mdp/actions.py | 4 +- .../locomanip_pick_place/mdp/actions.py | 7 +- .../mdp/pre_trained_policy_action.py | 4 +- .../contrib/ur10_particle_push/mdp/actions.py | 4 +- 28 files changed, 185 insertions(+), 234 deletions(-) delete mode 100644 source/isaaclab_contrib/changelog.d/mhaiderbhai-zero-actions.minor.rst delete mode 100644 source/isaaclab_newton/changelog.d/mhaiderbhai-zero-actions.minor.rst delete mode 100644 source/isaaclab_tasks/changelog.d/mhaiderbhai-zero-actions.minor.rst diff --git a/scripts/tutorials/03_envs/create_cube_base_env.py b/scripts/tutorials/03_envs/create_cube_base_env.py index 68ae34edb93a..43e269672344 100644 --- a/scripts/tutorials/03_envs/create_cube_base_env.py +++ b/scripts/tutorials/03_envs/create_cube_base_env.py @@ -121,9 +121,7 @@ def processed_actions(self) -> torch.Tensor: Operations """ - def process_actions(self, actions: torch.Tensor | None): - if actions is None: - actions = self._asset.data.root_pos_w.torch - self._env.scene.env_origins + def process_actions(self, actions: torch.Tensor): # store the raw actions self._raw_actions[:] = actions # no-processing of actions diff --git a/source/isaaclab/changelog.d/mhaiderbhai-neutral-actions.minor.rst b/source/isaaclab/changelog.d/mhaiderbhai-neutral-actions.minor.rst index 1bbf30477457..9343ebf809c6 100644 --- a/source/isaaclab/changelog.d/mhaiderbhai-neutral-actions.minor.rst +++ b/source/isaaclab/changelog.d/mhaiderbhai-neutral-actions.minor.rst @@ -1,5 +1,5 @@ Added ^^^^^ -* Allowed manager-based environments to accept ``None`` actions, with action terms applying their zero-action behavior. - Absolute differential IK, Pink IK, RMPFlow, and operational-space controllers hold their current poses. +* Added semantic neutral actions for manager-based action terms, including absolute differential IK, + Pink IK, RMPFlow, and operational-space controllers. diff --git a/source/isaaclab/isaaclab/envs/manager_based_env.py b/source/isaaclab/isaaclab/envs/manager_based_env.py index 38f2c3babde4..e427b7d374be 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_env.py +++ b/source/isaaclab/isaaclab/envs/manager_based_env.py @@ -500,7 +500,7 @@ def reset_to( # return observations return self.obs_buf, self.extras - def step(self, action: torch.Tensor | None) -> tuple[VecEnvObs, dict]: + def step(self, action: torch.Tensor) -> tuple[VecEnvObs, dict]: """Execute one time-step of the environment's dynamics. The environment steps forward at a fixed time-step, while the physics simulation is @@ -521,14 +521,13 @@ def step(self, action: torch.Tensor | None) -> tuple[VecEnvObs, dict]: app loop. Args: - action: The actions to apply on the environment. Shape is (num_envs, action_dim). If None, each action - term applies its zero-action behavior. + action: The actions to apply on the environment. Shape is (num_envs, action_dim). Returns: A tuple containing the observations and extras. """ # process actions - self.action_manager.process_action(action.to(self.device) if action is not None else None) + self.action_manager.process_action(action.to(self.device)) self.recorder_manager.record_pre_step() diff --git a/source/isaaclab/isaaclab/envs/manager_based_rl_env.py b/source/isaaclab/isaaclab/envs/manager_based_rl_env.py index 00e858e8ff81..ed751f1dd95e 100644 --- a/source/isaaclab/isaaclab/envs/manager_based_rl_env.py +++ b/source/isaaclab/isaaclab/envs/manager_based_rl_env.py @@ -173,7 +173,7 @@ def setup_manager_visualizers(self): Operations - MDP """ - def step(self, action: torch.Tensor | None) -> VecEnvStepReturn: + def step(self, action: torch.Tensor) -> VecEnvStepReturn: """Execute one time-step of the environment's dynamics and reset terminated environments. Unlike the :class:`ManagerBasedEnv.step` class, the function performs the following operations: @@ -199,14 +199,13 @@ def step(self, action: torch.Tensor | None) -> VecEnvStepReturn: - Post-reset re-renders for RTX sensors are also skipped. Args: - action: The actions to apply on the environment. Shape is (num_envs, action_dim). If None, each action - term applies its zero-action behavior. + action: The actions to apply on the environment. Shape is (num_envs, action_dim). Returns: A tuple containing the observations, rewards, resets (terminated and truncated) and extras. """ # process actions - self.action_manager.process_action(action.to(self.device) if action is not None else None) + self.action_manager.process_action(action.to(self.device)) self.recorder_manager.record_pre_step() diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/binary_joint_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/binary_joint_actions.py index 9b49ec77317e..db80bf2b6e64 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/binary_joint_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/binary_joint_actions.py @@ -129,9 +129,7 @@ def IO_descriptor(self) -> GenericActionIODescriptor: Operations. """ - def process_actions(self, actions: torch.Tensor | None): - if actions is None: - actions = self._raw_actions.zero_() + def process_actions(self, actions: torch.Tensor): # store the raw actions self._raw_actions[:] = actions # compute the binary mask @@ -184,9 +182,7 @@ class AbsBinaryJointPositionAction(BinaryJointAction): cfg: actions_cfg.AbsBinaryJointPositionActionCfg """The configuration of the action term.""" - def process_actions(self, actions: torch.Tensor | None): - if actions is None: - actions = self._raw_actions.zero_() + def process_actions(self, actions: torch.Tensor): # store the raw actions self._raw_actions[:] = actions # compute the binary mask diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py index 2d5a8cca4fe9..435a04574a06 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py @@ -167,9 +167,7 @@ def IO_descriptor(self) -> GenericActionIODescriptor: Operations. """ - def process_actions(self, actions: torch.Tensor | None): - if actions is None: - actions = self._raw_actions.zero_() + def process_actions(self, actions: torch.Tensor): # store the raw actions self._raw_actions[:] = actions # apply the affine transformations diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions_to_limits.py b/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions_to_limits.py index e6ba2867f987..f6c416b97337 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions_to_limits.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/joint_actions_to_limits.py @@ -152,9 +152,7 @@ def IO_descriptor(self) -> GenericActionIODescriptor: Operations. """ - def process_actions(self, actions: torch.Tensor | None): - if actions is None: - actions = self._raw_actions.zero_() + def process_actions(self, actions: torch.Tensor): # store the raw actions self._raw_actions[:] = actions # apply affine transformations @@ -281,7 +279,7 @@ def reset(self, env_ids: Sequence[int] | None = None) -> None: ) self._prev_applied_actions[env_ids, :] = curr_applied_actions - def process_actions(self, actions: torch.Tensor | None): + def process_actions(self, actions: torch.Tensor): # apply affine transformations super().process_actions(actions) # set position targets as moving average diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/non_holonomic_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/non_holonomic_actions.py index 7847493a5136..08d5e9be1c85 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/non_holonomic_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/non_holonomic_actions.py @@ -172,9 +172,7 @@ def IO_descriptor(self) -> GenericActionIODescriptor: Operations. """ - def process_actions(self, actions: torch.Tensor | None): - if actions is None: - actions = self._raw_actions.zero_() + def process_actions(self, actions): # store the raw actions self._raw_actions[:] = actions self._processed_actions = self.raw_actions * self._scale + self._offset diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py index 2590361d997b..32949738f918 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py @@ -170,6 +170,14 @@ def processed_actions(self) -> torch.Tensor: """Get the processed actions tensor.""" return self._processed_actions + @property + def neutral_actions(self) -> torch.Tensor: + """Raw actions that hold the controlled frames and hand joints at their current state.""" + frame_poses = self._asset.data.body_link_pose_w.torch[:, self._controlled_frame_ids].clone() + frame_poses[..., :3] -= self._env.scene.env_origins.unsqueeze(1) + hand_joint_positions = self._asset.data.joint_pos.torch[:, self._hand_joint_ids] + return torch.cat((frame_poses.flatten(start_dim=1), hand_joint_positions), dim=-1) + @property def IO_descriptor(self) -> GenericActionIODescriptor: """The IO descriptor of the action term. @@ -199,19 +207,12 @@ def IO_descriptor(self) -> GenericActionIODescriptor: # Operations. # """ - def process_actions(self, actions: torch.Tensor | None) -> None: + def process_actions(self, actions: torch.Tensor) -> None: """Process the input actions and set targets for each task. Args: - actions: The input actions tensor. If None, the current controlled-frame poses and hand-joint positions - are used. + actions: The input actions tensor. """ - if actions is None: - frame_poses = self._asset.data.body_link_pose_w.torch[:, self._controlled_frame_ids].clone() - frame_poses[..., :3] -= self._env.scene.env_origins.unsqueeze(1) - hand_joint_positions = self._asset.data.joint_pos.torch[:, self._hand_joint_ids] - actions = torch.cat((frame_poses.flatten(start_dim=1), hand_joint_positions), dim=-1) - # Store raw actions self._raw_actions[:] = actions diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/rmpflow_task_space_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/rmpflow_task_space_actions.py index 93a7fb09bbce..1893cc465f78 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/rmpflow_task_space_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/rmpflow_task_space_actions.py @@ -118,6 +118,16 @@ def raw_actions(self) -> torch.Tensor: def processed_actions(self) -> torch.Tensor: return self._processed_actions + @property + def neutral_actions(self) -> torch.Tensor: + """Raw actions that hold the current end-effector pose.""" + if self.cfg.use_relative_mode: + return super().neutral_actions + + ee_pos, ee_quat = self._compute_frame_pose() + command = torch.cat((ee_pos, ee_quat), dim=-1) + return torch.where(self._scale != 0.0, command / self._scale, torch.zeros_like(command)) + @property def jacobian_w(self) -> torch.Tensor: return self._asset.data.body_link_jacobian_w.torch[:, self._jacobi_body_idx, :, self._jacobi_joint_ids] @@ -136,14 +146,7 @@ def jacobian_b(self) -> torch.Tensor: """ # This is called each env.step() - def process_actions(self, actions: torch.Tensor | None): - if actions is None: - if self.cfg.use_relative_mode: - actions = self._raw_actions.zero_() - else: - ee_pos, ee_quat = self._compute_frame_pose() - command = torch.cat((ee_pos, ee_quat), dim=-1) - actions = torch.where(self._scale != 0.0, command / self._scale, torch.zeros_like(command)) + def process_actions(self, actions: torch.Tensor): # store the raw actions self._raw_actions[:] = actions self._processed_actions[:] = self.raw_actions * self._scale diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/surface_gripper_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/surface_gripper_actions.py index 62d4ff566cff..699743eb918b 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/surface_gripper_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/surface_gripper_actions.py @@ -85,9 +85,7 @@ def processed_actions(self) -> torch.Tensor: Operations. """ - def process_actions(self, actions: torch.Tensor | None): - if actions is None: - actions = self._raw_actions.zero_() + def process_actions(self, actions: torch.Tensor): # store the raw actions self._raw_actions[:] = actions # compute the binary mask diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py index a9e9f5c000dc..33adfba89b32 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py @@ -139,6 +139,16 @@ def raw_actions(self) -> torch.Tensor: def processed_actions(self) -> torch.Tensor: return self._processed_actions + @property + def neutral_actions(self) -> torch.Tensor: + """Raw actions that hold the current end-effector pose.""" + if self.cfg.controller.use_relative_mode: + return super().neutral_actions + + ee_pos, ee_quat = self._compute_frame_pose() + command = ee_pos if self.cfg.controller.command_type == "position" else torch.cat((ee_pos, ee_quat), dim=-1) + return torch.where(self._scale != 0.0, command / self._scale, torch.zeros_like(command)) + @property def jacobian_w(self) -> torch.Tensor: return self._asset.data.body_link_jacobian_w.torch[:, self._jacobi_body_idx, :, self._jacobi_joint_ids] @@ -187,16 +197,7 @@ def IO_descriptor(self) -> GenericActionIODescriptor: Operations. """ - def process_actions(self, actions: torch.Tensor | None): - if actions is None: - if self.cfg.controller.use_relative_mode: - actions = self._raw_actions.zero_() - else: - ee_pos, ee_quat = self._compute_frame_pose() - command = ( - ee_pos if self.cfg.controller.command_type == "position" else torch.cat((ee_pos, ee_quat), dim=-1) - ) - actions = torch.where(self._scale != 0.0, command / self._scale, torch.zeros_like(command)) + def process_actions(self, actions: torch.Tensor): # store the raw actions self._raw_actions[:] = actions self._processed_actions[:] = self.raw_actions * self._scale @@ -459,6 +460,40 @@ def processed_actions(self) -> torch.Tensor: """Processed actions for operational space control.""" return self._processed_actions + @property + def neutral_actions(self) -> torch.Tensor: + """Raw actions that hold the current end-effector pose and apply no wrench.""" + actions = super().neutral_actions + if self._pose_abs_idx is None: + return actions + + self._compute_ee_pose() + self._compute_task_frame_pose() + if self._task_frame_pose_b is None: + ee_pos_task = self._ee_pose_b[:, :3] + ee_quat_task = self._ee_pose_b[:, 3:7] + else: + ee_pos_task, ee_quat_task = math_utils.subtract_frame_transforms( + self._task_frame_pose_b[:, :3], + self._task_frame_pose_b[:, 3:7], + self._ee_pose_b[:, :3], + self._ee_pose_b[:, 3:7], + ) + + position_slice = slice(self._pose_abs_idx, self._pose_abs_idx + 3) + orientation_slice = slice(self._pose_abs_idx + 3, self._pose_abs_idx + 7) + actions[:, position_slice] = torch.where( + self._position_scale != 0.0, + ee_pos_task / self._position_scale, + torch.zeros_like(ee_pos_task), + ) + actions[:, orientation_slice] = torch.where( + self._orientation_scale != 0.0, + ee_quat_task / self._orientation_scale, + torch.zeros_like(ee_quat_task), + ) + return actions + @property def jacobian_w(self) -> torch.Tensor: return self._asset.data.body_link_jacobian_w.torch[:, self._jacobi_ee_body_idx, :, self._jacobi_joint_idx] @@ -517,13 +552,12 @@ def IO_descriptor(self) -> GenericActionIODescriptor: Operations. """ - def process_actions(self, actions: torch.Tensor | None): + def process_actions(self, actions: torch.Tensor): """Pre-processes the raw actions and sets them as commands for for operational space control. Args: - actions: The raw actions for operational space control. It is a tensor of shape - (``num_envs``, ``action_dim``). If None, the controller holds the current absolute pose and applies - zero relative pose and wrench commands. + actions (torch.Tensor): The raw actions for operational space control. It is a tensor of + shape (``num_envs``, ``action_dim``). """ # Update ee pose, which would be used by relative targets (i.e., pose_rel) @@ -532,30 +566,6 @@ def process_actions(self, actions: torch.Tensor | None): # Update task frame pose w.r.t. the root frame. self._compute_task_frame_pose() - if actions is None: - actions = self._raw_actions.zero_() - if self._pose_abs_idx is not None: - if self._task_frame_pose_b is None: - ee_pos_task = self._ee_pose_b[:, :3] - ee_quat_task = self._ee_pose_b[:, 3:7] - else: - ee_pos_task, ee_quat_task = math_utils.subtract_frame_transforms( - self._task_frame_pose_b[:, :3], - self._task_frame_pose_b[:, 3:7], - self._ee_pose_b[:, :3], - self._ee_pose_b[:, 3:7], - ) - actions[:, self._pose_abs_idx : self._pose_abs_idx + 3] = torch.where( - self._position_scale != 0.0, - ee_pos_task / self._position_scale, - torch.zeros_like(ee_pos_task), - ) - actions[:, self._pose_abs_idx + 3 : self._pose_abs_idx + 7] = torch.where( - self._orientation_scale != 0.0, - ee_quat_task / self._orientation_scale, - torch.zeros_like(ee_quat_task), - ) - # Pre-process the raw actions for operational space control. self._preprocess_actions(actions) diff --git a/source/isaaclab/isaaclab/managers/action_manager.py b/source/isaaclab/isaaclab/managers/action_manager.py index c58b5df70ba7..1e952628b168 100644 --- a/source/isaaclab/isaaclab/managers/action_manager.py +++ b/source/isaaclab/isaaclab/managers/action_manager.py @@ -88,6 +88,15 @@ def processed_actions(self) -> torch.Tensor: """The actions computed by the term after applying any processing.""" raise NotImplementedError + @property + def neutral_actions(self) -> torch.Tensor: + """Raw actions suitable for passive agent playback. + + The default is a zero-filled tensor. Action terms for which zero has a different or invalid meaning, + such as absolute-pose controllers, should override this property with a semantically neutral command. + """ + return torch.zeros_like(self.raw_actions) + @property def has_debug_vis_implementation(self) -> bool: """Whether the action term has a debug visualization implemented.""" @@ -139,14 +148,14 @@ def set_debug_vis(self, debug_vis: bool) -> bool: return True @abstractmethod - def process_actions(self, actions: torch.Tensor | None): + def process_actions(self, actions: torch.Tensor): """Processes the actions sent to the environment. Note: This function is called once per environment step by the manager. Args: - actions: The actions to process. If None, the action term applies its zero-action behavior. + actions: The actions to process. """ raise NotImplementedError @@ -263,6 +272,18 @@ def prev_action(self) -> torch.Tensor: """The previous actions sent to the environment. Shape is (num_envs, total_action_dim).""" return self._prev_action + @property + def neutral_actions(self) -> torch.Tensor: + """Raw actions suitable for passive playback of all active action terms. + + The returned tensor has shape ``(num_envs, total_action_dim)``. Since + some terms derive their neutral command from the current simulation + state, consumers should retrieve this property immediately before use. + """ + if not self._terms: + return torch.zeros_like(self._action) + return torch.cat([term.neutral_actions for term in self._terms.values()], dim=-1) + @property def has_debug_vis_implementation(self) -> bool: """Whether the command terms have debug visualization implemented.""" @@ -363,31 +384,26 @@ def reset(self, env_ids: Sequence[int] | None = None) -> dict[str, torch.Tensor] # nothing to log here return {} - def process_action(self, action: torch.Tensor | None): + def process_action(self, action: torch.Tensor): """Processes the actions sent to the environment. Note: This function should be called once per environment step. Args: - action: The actions to process. If None, each action term applies its zero-action behavior. + action: The actions to process. """ + # check if action dimension is valid + if self.total_action_dim != action.shape[1]: + raise ValueError(f"Invalid action shape, expected: {self.total_action_dim}, received: {action.shape[1]}.") + # store the input actions self._prev_action[:] = self._action - if action is None: - self._action.zero_() - else: - # check if action dimension is valid - if self.total_action_dim != action.shape[1]: - raise ValueError( - f"Invalid action shape, expected: {self.total_action_dim}, received: {action.shape[1]}." - ) - # store the input actions - self._action[:] = action.to(self.device) + self._action[:] = action.to(self.device) # split the actions and apply to each tensor idx = 0 for term in self._terms.values(): - term_actions = None if action is None else self._action[:, idx : idx + term.action_dim] + term_actions = action[:, idx : idx + term.action_dim] term.process_actions(term_actions) idx += term.action_dim diff --git a/source/isaaclab/test/envs/check_manager_based_env_floating_cube.py b/source/isaaclab/test/envs/check_manager_based_env_floating_cube.py index efb22d327dd9..ef0a151434b5 100644 --- a/source/isaaclab/test/envs/check_manager_based_env_floating_cube.py +++ b/source/isaaclab/test/envs/check_manager_based_env_floating_cube.py @@ -120,9 +120,7 @@ def processed_actions(self) -> torch.Tensor: Operations """ - def process_actions(self, actions: torch.Tensor | None): - if actions is None: - actions = self._asset.data.root_pos_w.torch - self._env.scene.env_origins + def process_actions(self, actions: torch.Tensor): # store the raw actions self._raw_actions[:] = actions # no-processing of actions diff --git a/source/isaaclab/test/envs/test_neutral_actions.py b/source/isaaclab/test/envs/test_neutral_actions.py index 2af8b5296b75..d0b6708c314d 100644 --- a/source/isaaclab/test/envs/test_neutral_actions.py +++ b/source/isaaclab/test/envs/test_neutral_actions.py @@ -3,44 +3,17 @@ # # SPDX-License-Identifier: BSD-3-Clause -"""Tests for action-term zero actions.""" +"""Tests for semantic neutral actions.""" from types import SimpleNamespace import torch from isaaclab.envs.mdp.actions.pink_task_space_actions import PinkInverseKinematicsAction -from isaaclab.managers.action_manager import ActionManager -def test_action_manager_dispatches_none_and_records_zero_action() -> None: - """The action manager lets terms resolve None while recording a conceptual zero action.""" - - class _ActionTerm: - action_dim = 2 - raw_actions = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) - received_none = False - - def process_actions(self, actions: torch.Tensor | None) -> None: - assert actions is None - self.received_none = True - - manager = object.__new__(ActionManager) - term = _ActionTerm() - manager._terms = {"term": term} - manager._action = torch.full((2, 2), 5.0) - manager._prev_action = torch.full((2, 2), -1.0) - manager._resolve_terms_handle = None - - manager.process_action(None) - - assert term.received_none - assert torch.equal(manager.action, torch.zeros(2, 2)) - assert torch.equal(manager.prev_action, torch.full((2, 2), 5.0)) - - -def test_pink_none_actions_use_current_frame_poses() -> None: - """Pink IK resolves None to valid current poses and hand joint positions.""" +def test_pink_neutral_actions_use_current_frame_poses() -> None: + """Pink IK neutral actions contain valid current poses and hand joint positions.""" action_term = object.__new__(PinkInverseKinematicsAction) action_term._controlled_frame_ids = [1, 0] action_term._hand_joint_ids = [1, 3] @@ -60,19 +33,11 @@ def test_pink_none_actions_use_current_frame_poses() -> None: ) ) action_term._env = SimpleNamespace(scene=SimpleNamespace(env_origins=env_origins)) - action_term.cfg = SimpleNamespace(controller=SimpleNamespace(num_hand_joints=2)) - action_term._raw_actions = torch.zeros(2, 16) - action_term._get_base_link_frame_transform = lambda: torch.eye(4).repeat(2, 1, 1) - action_term._extract_controlled_frame_poses = lambda actions: actions[:, :14] - action_term._transform_poses_to_base_link_frame = lambda poses: poses - action_term._set_task_targets = lambda poses: None - action_term.process_actions(None) + actions = action_term.neutral_actions expected_poses = body_poses[:, [1, 0]].clone() expected_poses[..., :3] -= env_origins.unsqueeze(1) expected = torch.cat((expected_poses.flatten(start_dim=1), joint_positions[:, [1, 3]]), dim=-1) - assert torch.equal(action_term.raw_actions, expected) - assert torch.all( - torch.linalg.vector_norm(action_term.raw_actions[:, :14].reshape(2, 2, 7)[..., 3:7], dim=-1) == 1.0 - ) + assert torch.equal(actions, expected) + assert torch.all(torch.linalg.vector_norm(actions[:, :14].reshape(2, 2, 7)[..., 3:7], dim=-1) == 1.0) diff --git a/source/isaaclab/test/envs/test_scale_randomization.py b/source/isaaclab/test/envs/test_scale_randomization.py index c007e754eb1f..d88029a8965e 100644 --- a/source/isaaclab/test/envs/test_scale_randomization.py +++ b/source/isaaclab/test/envs/test_scale_randomization.py @@ -95,9 +95,7 @@ def processed_actions(self) -> torch.Tensor: Operations """ - def process_actions(self, actions: torch.Tensor | None): - if actions is None: - actions = self._asset.data.root_pos_w.torch - self._env.scene.env_origins + def process_actions(self, actions: torch.Tensor): # store the raw actions self._raw_actions[:] = actions # no-processing of actions diff --git a/source/isaaclab_contrib/changelog.d/mhaiderbhai-zero-actions.minor.rst b/source/isaaclab_contrib/changelog.d/mhaiderbhai-zero-actions.minor.rst deleted file mode 100644 index 715e429a6a9d..000000000000 --- a/source/isaaclab_contrib/changelog.d/mhaiderbhai-zero-actions.minor.rst +++ /dev/null @@ -1,4 +0,0 @@ -Added -^^^^^ - -* Added support for ``None`` zero actions to multirotor thrust and navigation action terms. diff --git a/source/isaaclab_contrib/isaaclab_contrib/mdp/actions/thrust_actions.py b/source/isaaclab_contrib/isaaclab_contrib/mdp/actions/thrust_actions.py index 852b765641fe..897d621246b9 100644 --- a/source/isaaclab_contrib/isaaclab_contrib/mdp/actions/thrust_actions.py +++ b/source/isaaclab_contrib/isaaclab_contrib/mdp/actions/thrust_actions.py @@ -200,7 +200,7 @@ def reset(self, env_ids: Sequence[int] | None = None) -> None: """ self._raw_actions[env_ids] = 0.0 - def process_actions(self, actions: torch.Tensor | None): + def process_actions(self, actions: torch.Tensor): r"""Process actions by applying scaling, offset, and clipping. This method transforms raw policy actions into thrust commands through @@ -216,14 +216,12 @@ def process_actions(self, actions: torch.Tensor | None): Args: actions: Raw action tensor from the policy. Shape is ``(num_envs, action_dim)``. - Typically in the range [-1, 1] for normalized policies. If None, zeros are used. + Typically in the range [-1, 1] for normalized policies. Note: The processed actions are stored internally and applied during the next :meth:`apply_actions` call. """ - if actions is None: - actions = self._raw_actions.zero_() # store the raw actions self._raw_actions[:] = actions # apply the affine transformations @@ -333,7 +331,7 @@ def IO_descriptor(self) -> GenericActionIODescriptor: descriptor.action_type = "NavigationAction" return descriptor - def process_actions(self, actions: torch.Tensor | None): + def process_actions(self, actions: torch.Tensor): """Process actions by applying scaling, offset, and clipping.""" # Call parent to handle basic processing super().process_actions(actions) diff --git a/source/isaaclab_newton/changelog.d/mhaiderbhai-zero-actions.minor.rst b/source/isaaclab_newton/changelog.d/mhaiderbhai-zero-actions.minor.rst deleted file mode 100644 index a548ecbd5554..000000000000 --- a/source/isaaclab_newton/changelog.d/mhaiderbhai-zero-actions.minor.rst +++ /dev/null @@ -1,4 +0,0 @@ -Added -^^^^^ - -* Added support for ``None`` zero actions to Newton inverse-kinematics action terms, which hold their current poses. diff --git a/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py b/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py index 5fc3aed004d2..498aef87ab19 100644 --- a/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py +++ b/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_ik_actions.py @@ -50,7 +50,6 @@ def _ik_world_target_kernel( scale: wp.array(dtype=wp.float32), command_code: int, use_relative: int, - zero_action: int, out_pos: wp.array(dtype=wp.vec3f), out_rot: wp.array(dtype=wp.vec4f), ): @@ -72,43 +71,42 @@ def _ik_world_target_kernel( target_pos = ee_pos target_rot = ee_rot - if zero_action == 0: - if command_code == 0: # COMMAND_POSITION - disp = wp.vec3f( + if command_code == 0: # COMMAND_POSITION + disp = wp.vec3f( + action[i, action_offset + 0] * scale[0], + action[i, action_offset + 1] * scale[1], + action[i, action_offset + 2] * scale[2], + ) + target_pos = ee_pos + disp if use_relative == 1 else disp + else: + if use_relative == 1: + target_pos = ee_pos + wp.vec3f( action[i, action_offset + 0] * scale[0], action[i, action_offset + 1] * scale[1], action[i, action_offset + 2] * scale[2], ) - target_pos = ee_pos + disp if use_relative == 1 else disp + rot_vec = wp.vec3f( + action[i, action_offset + 3] * scale[3], + action[i, action_offset + 4] * scale[4], + action[i, action_offset + 5] * scale[5], + ) + angle = wp.length(rot_vec) + delta_rot = wp.quat_identity() + if angle > 1.0e-6: + delta_rot = wp.quat_from_axis_angle(rot_vec / angle, angle) + target_rot = delta_rot * ee_rot else: - if use_relative == 1: - target_pos = ee_pos + wp.vec3f( - action[i, action_offset + 0] * scale[0], - action[i, action_offset + 1] * scale[1], - action[i, action_offset + 2] * scale[2], - ) - rot_vec = wp.vec3f( - action[i, action_offset + 3] * scale[3], - action[i, action_offset + 4] * scale[4], - action[i, action_offset + 5] * scale[5], - ) - angle = wp.length(rot_vec) - delta_rot = wp.quat_identity() - if angle > 1.0e-6: - delta_rot = wp.quat_from_axis_angle(rot_vec / angle, angle) - target_rot = delta_rot * ee_rot - else: - target_pos = wp.vec3f( - action[i, action_offset + 0] * scale[0], - action[i, action_offset + 1] * scale[1], - action[i, action_offset + 2] * scale[2], - ) - target_rot = wp.quatf( - action[i, action_offset + 3] * scale[3], - action[i, action_offset + 4] * scale[4], - action[i, action_offset + 5] * scale[5], - action[i, action_offset + 6] * scale[6], - ) + target_pos = wp.vec3f( + action[i, action_offset + 0] * scale[0], + action[i, action_offset + 1] * scale[1], + action[i, action_offset + 2] * scale[2], + ) + target_rot = wp.quatf( + action[i, action_offset + 3] * scale[3], + action[i, action_offset + 4] * scale[4], + action[i, action_offset + 5] * scale[5], + action[i, action_offset + 6] * scale[6], + ) # Broadcast against the env-0 prototype root (all roots identical, validated). world_t = wp.transform_multiply(wp.transformf(root_pos_w[0], root_quat_w[0]), wp.transformf(target_pos, target_rot)) @@ -268,10 +266,7 @@ def IO_descriptor(self) -> GenericActionIODescriptor: self._IO_descriptor.extras["coordinate_names"] = self._action_coordinate_names() return self._IO_descriptor - def process_actions(self, actions: torch.Tensor | None) -> None: - zero_action = actions is None - if actions is None: - actions = self._raw_actions.zero_() + def process_actions(self, actions: torch.Tensor) -> None: self._raw_actions[:] = actions self._processed_actions[:] = self._raw_actions if self._clip is not None: @@ -303,7 +298,6 @@ def process_actions(self, actions: torch.Tensor | None) -> None: obj.scale, obj.command_code, obj.use_relative, - int(zero_action), obj.position_objective.target_positions, obj.rotation_objective.target_rotations, ], diff --git a/source/isaaclab_rl/changelog.d/mhaiderbhai-zero-agent-actions.rst b/source/isaaclab_rl/changelog.d/mhaiderbhai-zero-agent-actions.rst index 68fb064e2a90..bfae95cb6773 100644 --- a/source/isaaclab_rl/changelog.d/mhaiderbhai-zero-agent-actions.rst +++ b/source/isaaclab_rl/changelog.d/mhaiderbhai-zero-agent-actions.rst @@ -1,5 +1,5 @@ Fixed ^^^^^ -* Fixed the zero agent to defer zero-action handling to manager-based action terms, support composite and multi-agent - action spaces, and reject invalid task configurations before launching the simulator. +* Fixed the zero agent to use semantic neutral actions, support composite and multi-agent action spaces, + and reject invalid task configurations before launching the simulator. diff --git a/source/isaaclab_rl/isaaclab_rl/entrypoints/simple_agents.py b/source/isaaclab_rl/isaaclab_rl/entrypoints/simple_agents.py index 7e91ff5bf9b2..b1b90c4717f7 100644 --- a/source/isaaclab_rl/isaaclab_rl/entrypoints/simple_agents.py +++ b/source/isaaclab_rl/isaaclab_rl/entrypoints/simple_agents.py @@ -6,7 +6,7 @@ """Checkpoint-free playback workflows for Isaac Lab environments. The zero and random agents are variations of playback that need no trained checkpoint: -the policy either requests each action term's zero-action behavior or samples uniform random actions. +the policy either emits neutral actions or samples uniform random actions. """ from __future__ import annotations @@ -47,7 +47,7 @@ def run(argv: list[str] | None = None, *, policy: PolicyName) -> None: Args: argv: Command-line arguments excluding the executable name. Reads ``sys.argv`` when omitted. - policy: Action policy to apply, either term-specific zero actions or uniform random actions. + policy: Action policy to apply, either neutral actions or uniform random actions. Raises: ValueError: If the requested policy is not supported. @@ -94,7 +94,7 @@ def run(argv: list[str] | None = None, *, policy: PolicyName) -> None: # run everything in inference mode with torch.inference_mode(): if policy == "zero": - actions = _get_zero_actions(env) + actions = _get_neutral_actions(env) else: # sample actions from -1 to 1 actions = 2 * torch.rand(env.action_space.shape, device=device) - 1 @@ -104,16 +104,18 @@ def run(argv: list[str] | None = None, *, policy: PolicyName) -> None: env.close() -def _get_zero_actions(env: gym.Env): - """Create zero actions for passive environment playback. +def _get_neutral_actions(env: gym.Env): + """Create semantically neutral actions for passive environment playback. - Manager-based environments accept None so that each action term can apply its zero-action behavior. Direct-workflow - environments use zero-filled samples of their declared Gymnasium spaces, including composite and multi-agent spaces. + Manager-based environments can provide semantic neutral actions for terms + where literal zeros are unsafe, such as absolute-pose IK. Direct-workflow + environments fall back to zero-filled samples of their declared Gymnasium + spaces, including composite and multi-agent spaces. """ unwrapped = env.unwrapped action_manager = getattr(unwrapped, "action_manager", None) if action_manager is not None: - return None + return action_manager.neutral_actions if hasattr(unwrapped, "action_spaces"): return { diff --git a/source/isaaclab_rl/test/test_entrypoints.py b/source/isaaclab_rl/test/test_entrypoints.py index 826f2f161986..a0ceb98525a8 100644 --- a/source/isaaclab_rl/test/test_entrypoints.py +++ b/source/isaaclab_rl/test/test_entrypoints.py @@ -20,14 +20,15 @@ from isaaclab_rl.entrypoints import PlaybackRequest, TrainingRequest, api, dispatch from isaaclab_rl.entrypoints import simple_agents as _simple_agents -from isaaclab_rl.entrypoints.simple_agents import _get_zero_actions +from isaaclab_rl.entrypoints.simple_agents import _get_neutral_actions -def test_zero_agent_defers_to_manager_action_terms() -> None: - """The zero agent lets manager-based action terms define their zero-action behavior.""" - unwrapped = SimpleNamespace(action_manager=SimpleNamespace()) +def test_zero_agent_uses_manager_semantic_neutral_actions() -> None: + """The zero agent honors action-term neutral commands instead of forcing literal zeros.""" + expected = torch.tensor([[0.1, 0.2, 0.3, 1.0]]) + unwrapped = SimpleNamespace(action_manager=SimpleNamespace(neutral_actions=expected)) - assert _get_zero_actions(SimpleNamespace(unwrapped=unwrapped)) is None + assert _get_neutral_actions(SimpleNamespace(unwrapped=unwrapped)) is expected def test_zero_agent_supports_composite_direct_action_spaces() -> None: @@ -45,7 +46,7 @@ def test_zero_agent_supports_composite_direct_action_spaces() -> None: num_envs=2, ) - actions = _get_zero_actions(SimpleNamespace(unwrapped=unwrapped)) + actions = _get_neutral_actions(SimpleNamespace(unwrapped=unwrapped)) assert torch.equal(actions["continuous"], torch.zeros(2, 2)) assert torch.equal(actions["discrete"], torch.zeros(2, 1, dtype=torch.int64)) @@ -63,7 +64,7 @@ def test_zero_agent_supports_direct_multi_agent_action_spaces() -> None: num_envs=3, ) - actions = _get_zero_actions(SimpleNamespace(unwrapped=unwrapped)) + actions = _get_neutral_actions(SimpleNamespace(unwrapped=unwrapped)) assert torch.equal(actions["robot"], torch.zeros(3, 2)) assert torch.equal(actions["object"], torch.zeros(3, 1, dtype=torch.int64)) diff --git a/source/isaaclab_tasks/changelog.d/mhaiderbhai-zero-actions.minor.rst b/source/isaaclab_tasks/changelog.d/mhaiderbhai-zero-actions.minor.rst deleted file mode 100644 index ee3dd361456d..000000000000 --- a/source/isaaclab_tasks/changelog.d/mhaiderbhai-zero-actions.minor.rst +++ /dev/null @@ -1,4 +0,0 @@ -Added -^^^^^ - -* Added support for ``None`` zero actions to task-specific action terms. diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/actions.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/actions.py index 5d44a42ec763..4e2491b4832a 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/actions.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/franka_pour/mdp/actions.py @@ -39,7 +39,7 @@ def __init__(self, cfg: EMARelativeJointPositionActionCfg, env: ManagerBasedEnv) raise ValueError(f"Moving-average weight must lie in (0, 1], got {self._alpha}.") self._previous_delta = torch.zeros_like(self._processed_actions) - def process_actions(self, actions: torch.Tensor | None) -> None: + def process_actions(self, actions: torch.Tensor) -> None: """Affine-map the raw action, then smooth only the commanded joint delta.""" super().process_actions(actions) self._processed_actions.lerp_(self._previous_delta, 1.0 - self._alpha) @@ -136,7 +136,7 @@ def set_reset_position( expanded = position.expand(-1, self._num_joints) self._processed_actions[selected] = expanded - def process_actions(self, actions: torch.Tensor | None) -> None: + def process_actions(self, actions: torch.Tensor) -> None: previous_target = self._processed_actions super().process_actions(actions) previous_target.lerp_(self._processed_actions, self._alpha) diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/locomanip_pick_place/mdp/actions.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/locomanip_pick_place/mdp/actions.py index 70678632ec3d..87996ca9b066 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/locomanip_pick_place/mdp/actions.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/locomanip_pick_place/mdp/actions.py @@ -92,16 +92,13 @@ def _compose_policy_input(self, base_command: torch.Tensor, obs_tensor: torch.Te return policy_input - def process_actions(self, actions: torch.Tensor | None): + def process_actions(self, actions: torch.Tensor): """Process the input actions using the locomotion policy. Args: - actions: The lower body commands. If None, zeros are used. + actions: The lower body commands. """ - if actions is None: - actions = torch.zeros((self.num_envs, self.action_dim), device=self.device) - # Extract base command from the action tensor # Assuming the base command [vx, vy, wz, hip_height] base_command = actions diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/navigation/mdp/pre_trained_policy_action.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/navigation/mdp/pre_trained_policy_action.py index 61073a9dc002..4857d63711e1 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/navigation/mdp/pre_trained_policy_action.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/navigation/mdp/pre_trained_policy_action.py @@ -90,9 +90,7 @@ def processed_actions(self) -> torch.Tensor: Operations. """ - def process_actions(self, actions: torch.Tensor | None): - if actions is None: - actions = self._raw_actions.zero_() + def process_actions(self, actions: torch.Tensor): self._raw_actions[:] = actions def apply_actions(self): diff --git a/source/isaaclab_tasks/isaaclab_tasks/contrib/ur10_particle_push/mdp/actions.py b/source/isaaclab_tasks/isaaclab_tasks/contrib/ur10_particle_push/mdp/actions.py index 7060053ae900..a8f01c5d6220 100644 --- a/source/isaaclab_tasks/isaaclab_tasks/contrib/ur10_particle_push/mdp/actions.py +++ b/source/isaaclab_tasks/isaaclab_tasks/contrib/ur10_particle_push/mdp/actions.py @@ -50,10 +50,8 @@ def invalid_actions(self) -> torch.Tensor: """Whether the latest policy action contained a non-finite component.""" return self._invalid_actions - def process_actions(self, actions: torch.Tensor | None) -> None: + def process_actions(self, actions: torch.Tensor) -> None: """Sanitize the policy action and construct one bounded joint target.""" - if actions is None: - actions = torch.zeros_like(self._raw_actions) self._previous_actions.copy_(self._raw_actions) self._invalid_actions.copy_(~torch.isfinite(actions).all(dim=1)) self._raw_actions.copy_(torch.nan_to_num(actions, nan=0.0, posinf=1.0, neginf=-1.0).clamp(-1.0, 1.0)) From ba567ae761ce061f118ea65cc999752f6502d535 Mon Sep 17 00:00:00 2001 From: Mustafa Haiderbhai Date: Mon, 31 Aug 2026 19:04:40 -0700 Subject: [PATCH 4/4] Infer safe zero-agent actions --- .../mhaiderbhai-neutral-actions.minor.rst | 5 - .../mdp/actions/pink_task_space_actions.py | 23 --- .../mdp/actions/rmpflow_task_space_actions.py | 10 -- .../envs/mdp/actions/task_space_actions.py | 44 ------ .../isaaclab/managers/action_manager.py | 21 --- .../test/envs/test_neutral_actions.py | 43 ------ .../mhaiderbhai-zero-agent-actions.rst | 4 +- .../isaaclab_rl/entrypoints/simple_agents.py | 134 ++++++++++++++++-- source/isaaclab_rl/test/test_entrypoints.py | 130 +++++++++++++++-- 9 files changed, 244 insertions(+), 170 deletions(-) delete mode 100644 source/isaaclab/changelog.d/mhaiderbhai-neutral-actions.minor.rst delete mode 100644 source/isaaclab/test/envs/test_neutral_actions.py diff --git a/source/isaaclab/changelog.d/mhaiderbhai-neutral-actions.minor.rst b/source/isaaclab/changelog.d/mhaiderbhai-neutral-actions.minor.rst deleted file mode 100644 index 9343ebf809c6..000000000000 --- a/source/isaaclab/changelog.d/mhaiderbhai-neutral-actions.minor.rst +++ /dev/null @@ -1,5 +0,0 @@ -Added -^^^^^ - -* Added semantic neutral actions for manager-based action terms, including absolute differential IK, - Pink IK, RMPFlow, and operational-space controllers. diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py index 32949738f918..3a5e67ffc426 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/pink_task_space_actions.py @@ -75,16 +75,6 @@ def _initialize_joint_info(self) -> None: # Resolve hand joints self._hand_joint_ids, self._hand_joint_names = self._asset.find_joints(self.cfg.hand_joint_names) - # Resolve controlled frames in the same order as their pose commands. - self._controlled_frame_ids, controlled_frame_names = self._asset.find_bodies( - list(self.cfg.target_eef_link_names.values()), preserve_order=True - ) - if len(self._controlled_frame_ids) != len(self.cfg.target_eef_link_names): - raise ValueError( - "Expected one controlled body for every Pink IK target. Resolved " - f"{controlled_frame_names} from {list(self.cfg.target_eef_link_names.values())}." - ) - # Combine all joint information self._controlled_joint_ids = self._isaaclab_controlled_joint_ids + self._hand_joint_ids self._controlled_joint_names = self._isaaclab_controlled_joint_names + self._hand_joint_names @@ -119,11 +109,6 @@ def _initialize_helper_tensors(self) -> None: 1 for task in self._ik_controllers[0].cfg.variable_input_tasks if isinstance(task, FrameTask) ) self._num_frame_tasks = num_frame_tasks - if len(self._controlled_frame_ids) != self._num_frame_tasks: - raise ValueError( - f"Pink IK has {self._num_frame_tasks} variable frame tasks but " - f"{len(self._controlled_frame_ids)} controlled bodies were configured." - ) self._controlled_frame_poses = torch.zeros(num_frame_tasks, self.num_envs, 4, 4, device=self.device) # Pre-allocate tensor for base frame computations @@ -170,14 +155,6 @@ def processed_actions(self) -> torch.Tensor: """Get the processed actions tensor.""" return self._processed_actions - @property - def neutral_actions(self) -> torch.Tensor: - """Raw actions that hold the controlled frames and hand joints at their current state.""" - frame_poses = self._asset.data.body_link_pose_w.torch[:, self._controlled_frame_ids].clone() - frame_poses[..., :3] -= self._env.scene.env_origins.unsqueeze(1) - hand_joint_positions = self._asset.data.joint_pos.torch[:, self._hand_joint_ids] - return torch.cat((frame_poses.flatten(start_dim=1), hand_joint_positions), dim=-1) - @property def IO_descriptor(self) -> GenericActionIODescriptor: """The IO descriptor of the action term. diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/rmpflow_task_space_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/rmpflow_task_space_actions.py index 1893cc465f78..6cccf308e1d5 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/rmpflow_task_space_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/rmpflow_task_space_actions.py @@ -118,16 +118,6 @@ def raw_actions(self) -> torch.Tensor: def processed_actions(self) -> torch.Tensor: return self._processed_actions - @property - def neutral_actions(self) -> torch.Tensor: - """Raw actions that hold the current end-effector pose.""" - if self.cfg.use_relative_mode: - return super().neutral_actions - - ee_pos, ee_quat = self._compute_frame_pose() - command = torch.cat((ee_pos, ee_quat), dim=-1) - return torch.where(self._scale != 0.0, command / self._scale, torch.zeros_like(command)) - @property def jacobian_w(self) -> torch.Tensor: return self._asset.data.body_link_jacobian_w.torch[:, self._jacobi_body_idx, :, self._jacobi_joint_ids] diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py index 33adfba89b32..ddc994ec52b0 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py +++ b/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py @@ -139,16 +139,6 @@ def raw_actions(self) -> torch.Tensor: def processed_actions(self) -> torch.Tensor: return self._processed_actions - @property - def neutral_actions(self) -> torch.Tensor: - """Raw actions that hold the current end-effector pose.""" - if self.cfg.controller.use_relative_mode: - return super().neutral_actions - - ee_pos, ee_quat = self._compute_frame_pose() - command = ee_pos if self.cfg.controller.command_type == "position" else torch.cat((ee_pos, ee_quat), dim=-1) - return torch.where(self._scale != 0.0, command / self._scale, torch.zeros_like(command)) - @property def jacobian_w(self) -> torch.Tensor: return self._asset.data.body_link_jacobian_w.torch[:, self._jacobi_body_idx, :, self._jacobi_joint_ids] @@ -460,40 +450,6 @@ def processed_actions(self) -> torch.Tensor: """Processed actions for operational space control.""" return self._processed_actions - @property - def neutral_actions(self) -> torch.Tensor: - """Raw actions that hold the current end-effector pose and apply no wrench.""" - actions = super().neutral_actions - if self._pose_abs_idx is None: - return actions - - self._compute_ee_pose() - self._compute_task_frame_pose() - if self._task_frame_pose_b is None: - ee_pos_task = self._ee_pose_b[:, :3] - ee_quat_task = self._ee_pose_b[:, 3:7] - else: - ee_pos_task, ee_quat_task = math_utils.subtract_frame_transforms( - self._task_frame_pose_b[:, :3], - self._task_frame_pose_b[:, 3:7], - self._ee_pose_b[:, :3], - self._ee_pose_b[:, 3:7], - ) - - position_slice = slice(self._pose_abs_idx, self._pose_abs_idx + 3) - orientation_slice = slice(self._pose_abs_idx + 3, self._pose_abs_idx + 7) - actions[:, position_slice] = torch.where( - self._position_scale != 0.0, - ee_pos_task / self._position_scale, - torch.zeros_like(ee_pos_task), - ) - actions[:, orientation_slice] = torch.where( - self._orientation_scale != 0.0, - ee_quat_task / self._orientation_scale, - torch.zeros_like(ee_quat_task), - ) - return actions - @property def jacobian_w(self) -> torch.Tensor: return self._asset.data.body_link_jacobian_w.torch[:, self._jacobi_ee_body_idx, :, self._jacobi_joint_idx] diff --git a/source/isaaclab/isaaclab/managers/action_manager.py b/source/isaaclab/isaaclab/managers/action_manager.py index 1e952628b168..d711596e5f5a 100644 --- a/source/isaaclab/isaaclab/managers/action_manager.py +++ b/source/isaaclab/isaaclab/managers/action_manager.py @@ -88,15 +88,6 @@ def processed_actions(self) -> torch.Tensor: """The actions computed by the term after applying any processing.""" raise NotImplementedError - @property - def neutral_actions(self) -> torch.Tensor: - """Raw actions suitable for passive agent playback. - - The default is a zero-filled tensor. Action terms for which zero has a different or invalid meaning, - such as absolute-pose controllers, should override this property with a semantically neutral command. - """ - return torch.zeros_like(self.raw_actions) - @property def has_debug_vis_implementation(self) -> bool: """Whether the action term has a debug visualization implemented.""" @@ -272,18 +263,6 @@ def prev_action(self) -> torch.Tensor: """The previous actions sent to the environment. Shape is (num_envs, total_action_dim).""" return self._prev_action - @property - def neutral_actions(self) -> torch.Tensor: - """Raw actions suitable for passive playback of all active action terms. - - The returned tensor has shape ``(num_envs, total_action_dim)``. Since - some terms derive their neutral command from the current simulation - state, consumers should retrieve this property immediately before use. - """ - if not self._terms: - return torch.zeros_like(self._action) - return torch.cat([term.neutral_actions for term in self._terms.values()], dim=-1) - @property def has_debug_vis_implementation(self) -> bool: """Whether the command terms have debug visualization implemented.""" diff --git a/source/isaaclab/test/envs/test_neutral_actions.py b/source/isaaclab/test/envs/test_neutral_actions.py deleted file mode 100644 index d0b6708c314d..000000000000 --- a/source/isaaclab/test/envs/test_neutral_actions.py +++ /dev/null @@ -1,43 +0,0 @@ -# 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 - -"""Tests for semantic neutral actions.""" - -from types import SimpleNamespace - -import torch - -from isaaclab.envs.mdp.actions.pink_task_space_actions import PinkInverseKinematicsAction - - -def test_pink_neutral_actions_use_current_frame_poses() -> None: - """Pink IK neutral actions contain valid current poses and hand joint positions.""" - action_term = object.__new__(PinkInverseKinematicsAction) - action_term._controlled_frame_ids = [1, 0] - action_term._hand_joint_ids = [1, 3] - - body_poses = torch.tensor( - [ - [[1.0, 2.0, 3.0, 0.0, 0.0, 0.0, 1.0], [4.0, 5.0, 6.0, 0.0, 0.0, 1.0, 0.0]], - [[7.0, 8.0, 9.0, 0.0, 1.0, 0.0, 0.0], [10.0, 11.0, 12.0, 1.0, 0.0, 0.0, 0.0]], - ] - ) - joint_positions = torch.tensor([[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8]]) - env_origins = torch.tensor([[1.0, 1.0, 1.0], [2.0, 2.0, 2.0]]) - action_term._asset = SimpleNamespace( - data=SimpleNamespace( - body_link_pose_w=SimpleNamespace(torch=body_poses), - joint_pos=SimpleNamespace(torch=joint_positions), - ) - ) - action_term._env = SimpleNamespace(scene=SimpleNamespace(env_origins=env_origins)) - - actions = action_term.neutral_actions - - expected_poses = body_poses[:, [1, 0]].clone() - expected_poses[..., :3] -= env_origins.unsqueeze(1) - expected = torch.cat((expected_poses.flatten(start_dim=1), joint_positions[:, [1, 3]]), dim=-1) - assert torch.equal(actions, expected) - assert torch.all(torch.linalg.vector_norm(actions[:, :14].reshape(2, 2, 7)[..., 3:7], dim=-1) == 1.0) diff --git a/source/isaaclab_rl/changelog.d/mhaiderbhai-zero-agent-actions.rst b/source/isaaclab_rl/changelog.d/mhaiderbhai-zero-agent-actions.rst index bfae95cb6773..6924d5048609 100644 --- a/source/isaaclab_rl/changelog.d/mhaiderbhai-zero-agent-actions.rst +++ b/source/isaaclab_rl/changelog.d/mhaiderbhai-zero-agent-actions.rst @@ -1,5 +1,5 @@ Fixed ^^^^^ -* Fixed the zero agent to use semantic neutral actions, support composite and multi-agent action spaces, - and reject invalid task configurations before launching the simulator. +* Fixed the zero agent to infer finite hold commands for absolute task-space controllers, support composite and + multi-agent action spaces, and reject invalid task configurations before launching the simulator. diff --git a/source/isaaclab_rl/isaaclab_rl/entrypoints/simple_agents.py b/source/isaaclab_rl/isaaclab_rl/entrypoints/simple_agents.py index b1b90c4717f7..848245f1099c 100644 --- a/source/isaaclab_rl/isaaclab_rl/entrypoints/simple_agents.py +++ b/source/isaaclab_rl/isaaclab_rl/entrypoints/simple_agents.py @@ -6,7 +6,7 @@ """Checkpoint-free playback workflows for Isaac Lab environments. The zero and random agents are variations of playback that need no trained checkpoint: -the policy either emits neutral actions or samples uniform random actions. +the policy either infers finite zero or hold actions or samples uniform random actions. """ from __future__ import annotations @@ -14,13 +14,15 @@ import argparse import contextlib import sys -from typing import Literal +from collections.abc import Callable +from typing import Any, Literal import gymnasium as gym import torch from isaaclab.app import add_launcher_args, launch_simulation from isaaclab.envs.utils.spaces import sample_space +from isaaclab.utils import math as math_utils import isaaclab_tasks # noqa: F401 from isaaclab_tasks.utils import ( @@ -47,7 +49,7 @@ def run(argv: list[str] | None = None, *, policy: PolicyName) -> None: Args: argv: Command-line arguments excluding the executable name. Reads ``sys.argv`` when omitted. - policy: Action policy to apply, either neutral actions or uniform random actions. + policy: Action policy to apply, either inferred zero actions or uniform random actions. Raises: ValueError: If the requested policy is not supported. @@ -82,6 +84,7 @@ def run(argv: list[str] | None = None, *, policy: PolicyName) -> None: print(f"[INFO]: Gym action space: {env.action_space}") # reset environment env.reset() + zero_action_policy = _create_zero_action_policy(env) if policy == "zero" else None # simulate environment # keep running while any visualizer is open, and until the step budget is exhausted sim = env.unwrapped.sim @@ -94,7 +97,7 @@ def run(argv: list[str] | None = None, *, policy: PolicyName) -> None: # run everything in inference mode with torch.inference_mode(): if policy == "zero": - actions = _get_neutral_actions(env) + actions = zero_action_policy() else: # sample actions from -1 to 1 actions = 2 * torch.rand(env.action_space.shape, device=device) - 1 @@ -104,26 +107,129 @@ def run(argv: list[str] | None = None, *, policy: PolicyName) -> None: env.close() -def _get_neutral_actions(env: gym.Env): - """Create semantically neutral actions for passive environment playback. +def _create_zero_action_policy(env: gym.Env) -> Callable[[], Any]: + """Create a policy that emits finite actions for passive environment playback. - Manager-based environments can provide semantic neutral actions for terms - where literal zeros are unsafe, such as absolute-pose IK. Direct-workflow - environments fall back to zero-filled samples of their declared Gymnasium - spaces, including composite and multi-agent spaces. + Manager-based environments infer hold commands for absolute task-space action terms and use literal zeros for all + other terms. Direct-workflow environments use zero-filled samples of their declared Gymnasium spaces, including + composite and multi-agent spaces. """ unwrapped = env.unwrapped action_manager = getattr(unwrapped, "action_manager", None) if action_manager is not None: - return action_manager.neutral_actions + return _create_manager_zero_action_policy(action_manager, unwrapped) if hasattr(unwrapped, "action_spaces"): - return { + actions = { agent: sample_space(space, unwrapped.device, batch_size=unwrapped.num_envs, fill_value=0) for agent, space in unwrapped.action_spaces.items() } - - return sample_space(unwrapped.single_action_space, unwrapped.device, batch_size=unwrapped.num_envs, fill_value=0) + return lambda: actions + + actions = sample_space(unwrapped.single_action_space, unwrapped.device, batch_size=unwrapped.num_envs, fill_value=0) + return lambda: actions + + +def _create_manager_zero_action_policy(action_manager: Any, env: Any) -> Callable[[], torch.Tensor]: + """Create a zero-action policy from the active action terms.""" + actions = torch.zeros_like(action_manager.action) + term_policies = [] + index = 0 + for term_name in action_manager.active_terms: + term = action_manager.get_term(term_name) + term_policy = _create_action_term_zero_policy(term, env) + if term_policy is not None: + term_policies.append((slice(index, index + term.action_dim), term_policy)) + index += term.action_dim + + def policy() -> torch.Tensor: + actions.zero_() + for action_slice, term_policy in term_policies: + actions[:, action_slice] = term_policy() + if not torch.isfinite(actions).all(): + raise RuntimeError("Zero agent inferred non-finite actions from the current environment state.") + return actions + + return policy + + +def _create_action_term_zero_policy(term: Any, env: Any) -> Callable[[], torch.Tensor] | None: + """Create the specialized zero-action policy required by an action term.""" + term_types = {cls.__name__ for cls in type(term).__mro__} + + if "PinkInverseKinematicsAction" in term_types: + controlled_frame_ids, controlled_frame_names = term._asset.find_bodies( + list(term.cfg.target_eef_link_names.values()), preserve_order=True + ) + if len(controlled_frame_ids) != len(term.cfg.target_eef_link_names): + raise ValueError( + "Expected one controlled body for every Pink IK target. Resolved " + f"{controlled_frame_names} from {list(term.cfg.target_eef_link_names.values())}." + ) + if len(controlled_frame_ids) != term._num_frame_tasks: + raise ValueError( + f"Pink IK has {term._num_frame_tasks} variable frame tasks but " + f"{len(controlled_frame_ids)} controlled bodies were configured." + ) + + def pink_policy() -> torch.Tensor: + frame_poses = term._asset.data.body_link_pose_w.torch[:, controlled_frame_ids].clone() + frame_poses[..., :3] -= env.scene.env_origins.unsqueeze(1) + hand_joint_positions = term._asset.data.joint_pos.torch[:, term._hand_joint_ids] + return torch.cat((frame_poses.flatten(start_dim=1), hand_joint_positions), dim=-1) + + return pink_policy + + if "DifferentialInverseKinematicsAction" in term_types and not term.cfg.controller.use_relative_mode: + + def differential_ik_policy() -> torch.Tensor: + ee_pos, ee_quat = term._compute_frame_pose() + command = ee_pos if term.cfg.controller.command_type == "position" else torch.cat((ee_pos, ee_quat), dim=-1) + return _unscale_action(command, term._scale) + + return differential_ik_policy + + if "RMPFlowAction" in term_types and not term.cfg.use_relative_mode: + + def rmpflow_policy() -> torch.Tensor: + ee_pos, ee_quat = term._compute_frame_pose() + return _unscale_action(torch.cat((ee_pos, ee_quat), dim=-1), term._scale) + + return rmpflow_policy + + if "OperationalSpaceControllerAction" in term_types and term._pose_abs_idx is not None: + term_actions = torch.zeros_like(term.raw_actions) + + def operational_space_policy() -> torch.Tensor: + term_actions.zero_() + term._compute_ee_pose() + term._compute_task_frame_pose() + if term._task_frame_pose_b is None: + ee_pos_task = term._ee_pose_b[:, :3] + ee_quat_task = term._ee_pose_b[:, 3:7] + else: + ee_pos_task, ee_quat_task = math_utils.subtract_frame_transforms( + term._task_frame_pose_b[:, :3], + term._task_frame_pose_b[:, 3:7], + term._ee_pose_b[:, :3], + term._ee_pose_b[:, 3:7], + ) + term_actions[:, term._pose_abs_idx : term._pose_abs_idx + 3] = _unscale_action( + ee_pos_task, term._position_scale + ) + term_actions[:, term._pose_abs_idx + 3 : term._pose_abs_idx + 7] = _unscale_action( + ee_quat_task, term._orientation_scale + ) + return term_actions + + return operational_space_policy + + return None + + +def _unscale_action(command: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + """Map a processed command back to policy-action coordinates without division by zero.""" + return torch.where(scale != 0.0, command / scale, torch.zeros_like(command)) def _parse_args(argv: list[str] | None, policy: PolicyName) -> argparse.Namespace: diff --git a/source/isaaclab_rl/test/test_entrypoints.py b/source/isaaclab_rl/test/test_entrypoints.py index a0ceb98525a8..40e3901a9ff3 100644 --- a/source/isaaclab_rl/test/test_entrypoints.py +++ b/source/isaaclab_rl/test/test_entrypoints.py @@ -20,15 +20,129 @@ from isaaclab_rl.entrypoints import PlaybackRequest, TrainingRequest, api, dispatch from isaaclab_rl.entrypoints import simple_agents as _simple_agents -from isaaclab_rl.entrypoints.simple_agents import _get_neutral_actions +from isaaclab_rl.entrypoints.simple_agents import _create_zero_action_policy + + +def test_zero_agent_infers_finite_manager_actions() -> None: + """The zero agent holds absolute task-space targets and zeros all other action terms.""" + + class DifferentialInverseKinematicsAction: + action_dim = 7 + cfg = SimpleNamespace(controller=SimpleNamespace(use_relative_mode=False, command_type="pose")) + _scale = torch.tensor([2.0, 2.0, 2.0, 1.0, 1.0, 1.0, 1.0]) + + def _compute_frame_pose(self): + return torch.tensor([[2.0, 4.0, 6.0]]), torch.tensor([[0.0, 0.0, 0.0, 1.0]]) + + class RMPFlowAction: + action_dim = 7 + cfg = SimpleNamespace(use_relative_mode=False) + _scale = torch.ones(7) + + def _compute_frame_pose(self): + return torch.tensor([[3.0, 2.0, 1.0]]), torch.tensor([[0.0, 0.0, 1.0, 0.0]]) + + class PinkInverseKinematicsAction: + action_dim = 16 + cfg = SimpleNamespace(target_eef_link_names={"left": "left_hand", "right": "right_hand"}) + _hand_joint_ids = [1, 3] + _num_frame_tasks = 2 + + def __init__(self): + body_poses = torch.tensor( + [ + [ + [1.0, 2.0, 3.0, 0.0, 0.0, 0.0, 1.0], + [4.0, 5.0, 6.0, 0.0, 0.0, 1.0, 0.0], + ] + ] + ) + self._asset = SimpleNamespace( + find_bodies=lambda expressions, preserve_order: ([1, 0], ["left_hand", "right_hand"]), + data=SimpleNamespace( + body_link_pose_w=SimpleNamespace(torch=body_poses), + joint_pos=SimpleNamespace(torch=torch.tensor([[0.1, 0.2, 0.3, 0.4]])), + ), + ) + + class OperationalSpaceControllerAction: + action_dim = 7 + raw_actions = torch.zeros(1, 7) + _pose_abs_idx = 0 + _position_scale = torch.ones(3) + _orientation_scale = torch.ones(4) + _task_frame_pose_b = None + _ee_pose_b = torch.tensor([[9.0, 8.0, 7.0, 0.0, 1.0, 0.0, 0.0]]) + + def _compute_ee_pose(self): + pass + + def _compute_task_frame_pose(self): + pass + + class JointAction: + action_dim = 2 + + terms = { + "diff_ik": DifferentialInverseKinematicsAction(), + "rmpflow": RMPFlowAction(), + "pink": PinkInverseKinematicsAction(), + "osc": OperationalSpaceControllerAction(), + "joints": JointAction(), + } + manager = SimpleNamespace( + action=torch.empty(1, sum(term.action_dim for term in terms.values())), + active_terms=list(terms), + get_term=terms.__getitem__, + ) + unwrapped = SimpleNamespace( + action_manager=manager, + scene=SimpleNamespace(env_origins=torch.tensor([[1.0, 1.0, 1.0]])), + ) + + actions = _create_zero_action_policy(SimpleNamespace(unwrapped=unwrapped))() + + expected_pink_poses = torch.tensor( + [[[3.0, 4.0, 5.0, 0.0, 0.0, 1.0, 0.0], [0.0, 1.0, 2.0, 0.0, 0.0, 0.0, 1.0]]] + ).flatten(start_dim=1) + expected = torch.cat( + ( + torch.tensor([[1.0, 2.0, 3.0, 0.0, 0.0, 0.0, 1.0]]), + torch.tensor([[3.0, 2.0, 1.0, 0.0, 0.0, 1.0, 0.0]]), + expected_pink_poses, + torch.tensor([[0.2, 0.4]]), + OperationalSpaceControllerAction._ee_pose_b, + torch.zeros(1, 2), + ), + dim=-1, + ) + + assert torch.equal(actions, expected) + assert torch.isfinite(actions).all() + +def test_zero_agent_rejects_non_finite_inferred_actions() -> None: + """The zero agent stops before a non-finite inferred action reaches the environment.""" -def test_zero_agent_uses_manager_semantic_neutral_actions() -> None: - """The zero agent honors action-term neutral commands instead of forcing literal zeros.""" - expected = torch.tensor([[0.1, 0.2, 0.3, 1.0]]) - unwrapped = SimpleNamespace(action_manager=SimpleNamespace(neutral_actions=expected)) + class DifferentialInverseKinematicsAction: + action_dim = 7 + cfg = SimpleNamespace(controller=SimpleNamespace(use_relative_mode=False, command_type="pose")) + _scale = torch.ones(7) + + def _compute_frame_pose(self): + return torch.full((1, 3), torch.nan), torch.tensor([[0.0, 0.0, 0.0, 1.0]]) + + term = DifferentialInverseKinematicsAction() + manager = SimpleNamespace( + action=torch.empty(1, term.action_dim), + active_terms=["ik"], + get_term=lambda name: term, + ) + unwrapped = SimpleNamespace(action_manager=manager) + policy = _create_zero_action_policy(SimpleNamespace(unwrapped=unwrapped)) - assert _get_neutral_actions(SimpleNamespace(unwrapped=unwrapped)) is expected + with pytest.raises(RuntimeError, match="inferred non-finite actions"): + policy() def test_zero_agent_supports_composite_direct_action_spaces() -> None: @@ -46,7 +160,7 @@ def test_zero_agent_supports_composite_direct_action_spaces() -> None: num_envs=2, ) - actions = _get_neutral_actions(SimpleNamespace(unwrapped=unwrapped)) + actions = _create_zero_action_policy(SimpleNamespace(unwrapped=unwrapped))() assert torch.equal(actions["continuous"], torch.zeros(2, 2)) assert torch.equal(actions["discrete"], torch.zeros(2, 1, dtype=torch.int64)) @@ -64,7 +178,7 @@ def test_zero_agent_supports_direct_multi_agent_action_spaces() -> None: num_envs=3, ) - actions = _get_neutral_actions(SimpleNamespace(unwrapped=unwrapped)) + actions = _create_zero_action_policy(SimpleNamespace(unwrapped=unwrapped))() assert torch.equal(actions["robot"], torch.zeros(3, 2)) assert torch.equal(actions["object"], torch.zeros(3, 1, dtype=torch.int64))