66"""Checkpoint-free playback workflows for Isaac Lab environments.
77
88The zero and random agents are variations of playback that need no trained checkpoint:
9- the policy either emits constant zero actions or samples uniform random actions.
9+ the policy either infers finite zero or hold actions or samples uniform random actions.
1010"""
1111
1212from __future__ import annotations
1313
1414import argparse
1515import contextlib
1616import sys
17- from typing import Literal
17+ from collections .abc import Callable
18+ from typing import Any , Literal
1819
1920import gymnasium as gym
2021import torch
2122
2223from isaaclab .app import add_launcher_args , launch_simulation
24+ from isaaclab .envs .utils .spaces import sample_space
25+ from isaaclab .utils import math as math_utils
2326
2427import isaaclab_tasks # noqa: F401
2528from isaaclab_tasks .utils import (
@@ -46,7 +49,7 @@ def run(argv: list[str] | None = None, *, policy: PolicyName) -> None:
4649
4750 Args:
4851 argv: Command-line arguments excluding the executable name. Reads ``sys.argv`` when omitted.
49- policy: Action policy to apply, either constant zero actions or uniform random actions.
52+ policy: Action policy to apply, either inferred zero actions or uniform random actions.
5053
5154 Raises:
5255 ValueError: If the requested policy is not supported.
@@ -61,13 +64,18 @@ def run(argv: list[str] | None = None, *, policy: PolicyName) -> None:
6164 # parse configuration via Hydra (supports preset selection, e.g. env.sim.physics=newton_mjwarp)
6265 env_cfg , _ = resolve_task_config (args_cli .task , "" )
6366
64- with launch_simulation (env_cfg , args_cli ):
65- # override with CLI arguments
66- env_cfg .scene .num_envs = args_cli .num_envs if args_cli .num_envs is not None else env_cfg .scene .num_envs
67- env_cfg .sim .device = args_cli .device if args_cli .device is not None else env_cfg .sim .device
68- if args_cli .disable_fabric :
69- env_cfg .sim .use_fabric = False
67+ # override with CLI arguments and reject unsupported configurations before
68+ # launching Kit or initializing a native physics backend.
69+ env_cfg .scene .num_envs = args_cli .num_envs if args_cli .num_envs is not None else env_cfg .scene .num_envs
70+ env_cfg .sim .device = args_cli .device if args_cli .device is not None else env_cfg .sim .device
71+ if args_cli .disable_fabric :
72+ env_cfg .sim .use_fabric = False
73+ try :
74+ env_cfg .validate ()
75+ except (TypeError , ValueError ) as exc :
76+ raise SystemExit (f"Invalid environment configuration: { exc } ) from None
7077
78+ with launch_simulation (env_cfg , args_cli ):
7179 # create environment
7280 env = gym .make (args_cli .task , cfg = env_cfg )
7381
@@ -76,11 +84,11 @@ def run(argv: list[str] | None = None, *, policy: PolicyName) -> None:
7684 print (f"[INFO]: Gym action space: { env .action_space } )
7785 # reset environment
7886 env .reset ()
87+ zero_action_policy = _create_zero_action_policy (env ) if policy == "zero" else None
7988 # simulate environment
8089 # keep running while any visualizer is open, and until the step budget is exhausted
8190 sim = env .unwrapped .sim
8291 device = env .unwrapped .device
83- zero_actions = torch .zeros (env .action_space .shape , device = device )
8492 step = 0
8593 while sim .is_headless_or_exist_active_visualizer ():
8694 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:
8997 # run everything in inference mode
9098 with torch .inference_mode ():
9199 if policy == "zero" :
92- actions = zero_actions
100+ actions = zero_action_policy ()
93101 else :
94102 # sample actions from -1 to 1
95103 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:
99107 env .close ()
100108
101109
110+ def _create_zero_action_policy (env : gym .Env ) -> Callable [[], Any ]:
111+ """Create a policy that emits finite actions for passive environment playback.
112+
113+ Manager-based environments infer hold commands for absolute task-space action terms and use literal zeros for all
114+ other terms. Direct-workflow environments use zero-filled samples of their declared Gymnasium spaces, including
115+ composite and multi-agent spaces.
116+ """
117+ unwrapped = env .unwrapped
118+ action_manager = getattr (unwrapped , "action_manager" , None )
119+ if action_manager is not None :
120+ return _create_manager_zero_action_policy (action_manager , unwrapped )
121+
122+ if hasattr (unwrapped , "action_spaces" ):
123+ actions = {
124+ agent : sample_space (space , unwrapped .device , batch_size = unwrapped .num_envs , fill_value = 0 )
125+ for agent , space in unwrapped .action_spaces .items ()
126+ }
127+ return lambda : actions
128+
129+ actions = sample_space (unwrapped .single_action_space , unwrapped .device , batch_size = unwrapped .num_envs , fill_value = 0 )
130+ return lambda : actions
131+
132+
133+ def _create_manager_zero_action_policy (action_manager : Any , env : Any ) -> Callable [[], torch .Tensor ]:
134+ """Create a zero-action policy from the active action terms."""
135+ actions = torch .zeros_like (action_manager .action )
136+ term_policies = []
137+ index = 0
138+ for term_name in action_manager .active_terms :
139+ term = action_manager .get_term (term_name )
140+ term_policy = _create_action_term_zero_policy (term , env )
141+ if term_policy is not None :
142+ term_policies .append ((slice (index , index + term .action_dim ), term_policy ))
143+ index += term .action_dim
144+
145+ def policy () -> torch .Tensor :
146+ actions .zero_ ()
147+ for action_slice , term_policy in term_policies :
148+ actions [:, action_slice ] = term_policy ()
149+ if not torch .isfinite (actions ).all ():
150+ raise RuntimeError ("Zero agent inferred non-finite actions from the current environment state." )
151+ return actions
152+
153+ return policy
154+
155+
156+ def _create_action_term_zero_policy (term : Any , env : Any ) -> Callable [[], torch .Tensor ] | None :
157+ """Create the specialized zero-action policy required by an action term."""
158+ term_types = {cls .__name__ for cls in type (term ).__mro__ }
159+
160+ if "PinkInverseKinematicsAction" in term_types :
161+ controlled_frame_ids , controlled_frame_names = term ._asset .find_bodies (
162+ list (term .cfg .target_eef_link_names .values ()), preserve_order = True
163+ )
164+ if len (controlled_frame_ids ) != len (term .cfg .target_eef_link_names ):
165+ raise ValueError (
166+ "Expected one controlled body for every Pink IK target. Resolved "
167+ f"{ controlled_frame_names } { list (term .cfg .target_eef_link_names .values ())}
168+ )
169+ if len (controlled_frame_ids ) != term ._num_frame_tasks :
170+ raise ValueError (
171+ f"Pink IK has { term ._num_frame_tasks }
172+ f"{ len (controlled_frame_ids )}
173+ )
174+
175+ def pink_policy () -> torch .Tensor :
176+ frame_poses = term ._asset .data .body_link_pose_w .torch [:, controlled_frame_ids ].clone ()
177+ frame_poses [..., :3 ] -= env .scene .env_origins .unsqueeze (1 )
178+ hand_joint_positions = term ._asset .data .joint_pos .torch [:, term ._hand_joint_ids ]
179+ return torch .cat ((frame_poses .flatten (start_dim = 1 ), hand_joint_positions ), dim = - 1 )
180+
181+ return pink_policy
182+
183+ if "DifferentialInverseKinematicsAction" in term_types and not term .cfg .controller .use_relative_mode :
184+
185+ def differential_ik_policy () -> torch .Tensor :
186+ ee_pos , ee_quat = term ._compute_frame_pose ()
187+ command = ee_pos if term .cfg .controller .command_type == "position" else torch .cat ((ee_pos , ee_quat ), dim = - 1 )
188+ return _unscale_action (command , term ._scale )
189+
190+ return differential_ik_policy
191+
192+ if "RMPFlowAction" in term_types and not term .cfg .use_relative_mode :
193+
194+ def rmpflow_policy () -> torch .Tensor :
195+ ee_pos , ee_quat = term ._compute_frame_pose ()
196+ return _unscale_action (torch .cat ((ee_pos , ee_quat ), dim = - 1 ), term ._scale )
197+
198+ return rmpflow_policy
199+
200+ if "OperationalSpaceControllerAction" in term_types and term ._pose_abs_idx is not None :
201+ term_actions = torch .zeros_like (term .raw_actions )
202+
203+ def operational_space_policy () -> torch .Tensor :
204+ term_actions .zero_ ()
205+ term ._compute_ee_pose ()
206+ term ._compute_task_frame_pose ()
207+ if term ._task_frame_pose_b is None :
208+ ee_pos_task = term ._ee_pose_b [:, :3 ]
209+ ee_quat_task = term ._ee_pose_b [:, 3 :7 ]
210+ else :
211+ ee_pos_task , ee_quat_task = math_utils .subtract_frame_transforms (
212+ term ._task_frame_pose_b [:, :3 ],
213+ term ._task_frame_pose_b [:, 3 :7 ],
214+ term ._ee_pose_b [:, :3 ],
215+ term ._ee_pose_b [:, 3 :7 ],
216+ )
217+ term_actions [:, term ._pose_abs_idx : term ._pose_abs_idx + 3 ] = _unscale_action (
218+ ee_pos_task , term ._position_scale
219+ )
220+ term_actions [:, term ._pose_abs_idx + 3 : term ._pose_abs_idx + 7 ] = _unscale_action (
221+ ee_quat_task , term ._orientation_scale
222+ )
223+ return term_actions
224+
225+ return operational_space_policy
226+
227+ return None
228+
229+
230+ def _unscale_action (command : torch .Tensor , scale : torch .Tensor ) -> torch .Tensor :
231+ """Map a processed command back to policy-action coordinates without division by zero."""
232+ return torch .where (scale != 0.0 , command / scale , torch .zeros_like (command ))
233+
234+
102235def _parse_args (argv : list [str ] | None , policy : PolicyName ) -> argparse .Namespace :
103236 """Parse the command line of a checkpoint-free agent and hand the remainder to Hydra."""
104237 parser = argparse .ArgumentParser (description = _DESCRIPTIONS [policy ])
0 commit comments