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..6924d5048609 --- /dev/null +++ b/source/isaaclab_rl/changelog.d/mhaiderbhai-zero-agent-actions.rst @@ -0,0 +1,5 @@ +Fixed +^^^^^ + +* 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 fea6da73474f..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 constant zero 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,12 +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 ( @@ -46,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 constant zero 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. @@ -61,13 +64,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) @@ -76,11 +84,11 @@ 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 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 +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 = zero_actions + actions = zero_action_policy() else: # sample actions from -1 to 1 actions = 2 * torch.rand(env.action_space.shape, device=device) - 1 @@ -99,6 +107,131 @@ def run(argv: list[str] | None = None, *, policy: PolicyName) -> None: env.close() +def _create_zero_action_policy(env: gym.Env) -> Callable[[], Any]: + """Create a policy that emits finite actions for passive environment playback. + + 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 _create_manager_zero_action_policy(action_manager, unwrapped) + + if hasattr(unwrapped, "action_spaces"): + actions = { + agent: sample_space(space, unwrapped.device, batch_size=unwrapped.num_envs, fill_value=0) + for agent, space in unwrapped.action_spaces.items() + } + 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: """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..40e3901a9ff3 100644 --- a/source/isaaclab_rl/test/test_entrypoints.py +++ b/source/isaaclab_rl/test/test_entrypoints.py @@ -11,11 +11,200 @@ 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 _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.""" + + 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)) + + with pytest.raises(RuntimeError, match="inferred non-finite actions"): + policy() + + +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 = _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)) + + +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 = _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)) + + +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: