from typing import TYPE_CHECKING # noqa: I001 import numpy as np from omegaconf import OmegaConf import omni from pxr import Gf, UsdGeom import torch from loguru import logger from gear_sonic.trl.utils.common import custom_instantiate if TYPE_CHECKING: from isaaclab.envs.manager_based_rl_env import ManagerBasedEnv # Import joint index functions (single source of truth) from gear_sonic.envs.env_utils.joint_utils import get_body_joint_indices, get_hand_joint_indices # Import visualization markers for contact point visualization try: from isaaclab.markers import VisualizationMarkers, VisualizationMarkersCfg import isaaclab.sim as sim_utils VISUALIZATION_AVAILABLE = True except ImportError: VISUALIZATION_AVAILABLE = False class ManagerEnvWrapper: def __init__(self, env: "ManagerBasedEnv", config): env.wrapper = self self.env = env self.config = OmegaConf.create(config) self.device = env.device self.viewer_focused = True self.is_manager_env = True if hasattr(self.env, "num_envs"): self.num_envs = self.env.num_envs else: self.num_envs = self.env.env.unwrapped.num_envs try: self.motion_command = env.command_manager.get_term("motion") self._motion_lib = self.motion_command.motion_lib except: # noqa: E722 logger.info("No motion lib found") self.motion_command = None self._motion_lib = None try: self.force_command = env.command_manager.get_term("force") except: # noqa: E722 self.force_command = None self.is_evaluating = False self.start_idx = 0 self._last_predicted_object_pos = ( None # For debug visualization of predicted object position ) if not self.config.headless: self.setup_keyboard() # Action visualization toggle and state self.turn_on_visualization = bool(self.config.get("turn_on_visualization", False)) and ( not self.config.get("headless", False) ) self._viz_every_n_steps = int(self.config.get("viz_every_n_steps", 1)) self._plot_action_dim = int(self.config.get("action_plot_dim", 29)) clip_default = self.config.get("action_clip_value", 1.0) # noqa: F841 self._action_ylim = float(self.config.get("action_plot_ylim", 10.0)) self._plot_window = int(self.config.get("action_plot_window", 200)) self._step_counter = 0 self._action_fig = None self._action_lines = None self._action_hist = None self._hist_idx = 0 self._blit_background = None self._blit_supported = True self._blit_refresh_interval = int(self.config.get("plot_blit_refresh_interval", 120)) # Initialize action transform module from config self.action_transform_module = None self._needs_policy_atm = False # Default: use policy obs directly self._policy_atm_indices = None # Initialize joint index mappings (for DOF mismatch handling in replay and step) self._body_joint_indices = None self._hand_joint_indices = None self._setup_replay_joint_indices() # Setup for replay mode # Initialize finger primitive support self._use_finger_primitive = self.config.get("use_finger_primitive", False) self._finger_primitive_map = None if self._use_finger_primitive: self._setup_finger_primitives() # Latent residual mode: policy outputs residual added to token latent space self._use_latent_residual = self.config.get("use_latent_residual", False) # Latent residual options (only used when use_latent_residual=True) self._latent_residual_mode = self.config.get("latent_residual_mode", "post_quantization") self._latent_residual_scale = self.config.get("latent_residual_scale", 1.0) # Student direct latent mode: policy outputs FULL latent (not residual) # This is used for vision student policies that learned to output full latent # directly, bypassing the ATM encoder entirely at inference time. # The 64-dim output goes directly to ATM decoder (no encoding step). self._use_student_direct_latent = self.config.get("use_student_direct_latent", False) if self._use_latent_residual: logger.info( f"Latent residual enabled: mode={self._latent_residual_mode}, " f"scale={self._latent_residual_scale}" ) if self._use_student_direct_latent: logger.info( "Student direct latent mode enabled: policy output goes directly to ATM decoder " "(no encoding step)" ) # Camera extrinsics randomization state self._camera_extrinsics_randomized = False self._camera_base_transforms = {} # Store original transforms per env action_transform_module_cfg = self.config.get("action_transform_module_cfg", None) if action_transform_module_cfg is not None: # Load configs from exported YAML file with open(action_transform_module_cfg) as f: exported_config = OmegaConf.load(f) env_config = exported_config.get("env_config", {}) algo_config = exported_config.get("algo_config", {}) self.action_transform_module = custom_instantiate( algo_config.actor, env_config=env_config, algo_config=algo_config, _resolve=False ).to(self.device) logger.info(f"Initialized action_transform_module from config: {action_transform_module_cfg}") # Load checkpoint if provided action_transform_module_checkpoint = self.config.get( "action_transform_module_checkpoint", None ) if action_transform_module_checkpoint is not None: # Compatibility shim: checkpoints saved with TRL < 0.28.0 reference # trl.trainer.utils.OnlineTrainerState, which was moved in 0.28.0 try: from trl.experimental.ppo.ppo_trainer import OnlineTrainerState, exact_div import trl.trainer.utils trl.trainer.utils.OnlineTrainerState = OnlineTrainerState trl.trainer.utils.exact_div = exact_div except ImportError: pass checkpoint = torch.load( action_transform_module_checkpoint, map_location=self.device, weights_only=False ) self.action_transform_module.load_state_dict(checkpoint["policy_state_dict"]) logger.info( f"Loaded action_transform_module checkpoint: {action_transform_module_checkpoint}" ) # Precompute tokenizer observation indices for meta_action target self._tokenizer_obs_indices = self._compute_tokenizer_obs_indices() # Setup policy_atm observations for action_transform_module if DOF mismatch exists self._setup_policy_atm(env_config, algo_config) try: self.viewer_focused = True self.env.viewport_camera_controller.update_view_to_world() self.env.viewport_camera_controller.update_view_to_asset_root("robot") except: # noqa: E722 self.viewer_focused = False def _setup_replay_joint_indices(self): """Setup joint indices for replay mode if robot has more DOFs than motion lib (29).""" if self.env.scene["robot"].num_joints > 29: self._setup_action_joint_indices() def _compute_tokenizer_obs_indices(self): """Compute start and end indices for each tokenizer observation in the flattened tensor.""" if self.action_transform_module is None: return {} tokenizer_obs_names = self.action_transform_module.actor_module.tokenizer_obs_names tokenizer_obs_dims = self.action_transform_module.actor_module.tokenizer_obs_dims indices = {} current_index = 0 for name in tokenizer_obs_names: all_dim = int(np.prod(tokenizer_obs_dims[name])) indices[name] = (current_index, current_index + all_dim) current_index += all_dim return indices def _setup_policy_atm(self, env_config, algo_config): # noqa: ARG002 """Setup policy_atm for action_transform_module when robot has more DOFs than ATM expects.""" atm_num_joints = env_config.get("robot", {}).get("actions_dim") or env_config.get( "robot", {} ).get("num_joints", 29) self._needs_policy_atm = self.config.get("needs_policy_atm", True) self._atm_num_joints = atm_num_joints self._current_num_joints = self.env.scene["robot"].num_joints if self._current_num_joints > atm_num_joints: assert ( self._needs_policy_atm ), "Robot has more DOFs than ATM expects, but needs_policy_atm is False" has_policy_atm = ( hasattr(self.env, "observation_manager") and "policy_atm" in self.env.observation_manager._group_obs_term_names # noqa: SLF001 ) self._use_policy_atm_group = self._needs_policy_atm and has_policy_atm if self._use_policy_atm_group: self._setup_action_joint_indices() def _setup_action_joint_indices(self): """Compute joint indices for mapping body (29 DOF) and hand (14 DOF) actions.""" if self._body_joint_indices is not None: return robot = self.env.scene["robot"] self._body_joint_indices = get_body_joint_indices(robot) self._hand_joint_indices = get_hand_joint_indices(robot) def _setup_finger_primitives(self): """Setup finger primitive action mapping from config. Finger primitives allow the policy to output 2 actions (left/right gripper) instead of 14 individual finger joint actions. Each primitive action interpolates between open (pos_0) and closed (pos_1) positions. """ primitive_cfg = self.config.get("finger_primitive", {}) primitive_action_map = primitive_cfg.get("primitive_action_map", {}) if not primitive_action_map: logger.info("Warning: use_finger_primitive=True but no primitive_action_map defined") self._use_finger_primitive = False return robot = self.env.scene["robot"] joint_names = robot.joint_names self._finger_primitive_map = [] for action_name in sorted(primitive_action_map.keys()): prim_cfg = primitive_action_map[action_name] dof_names = list(prim_cfg.get("dof_names", [])) pos_0 = torch.tensor(prim_cfg.get("pos_0", []), device=self.device, dtype=torch.float32) pos_1 = torch.tensor(prim_cfg.get("pos_1", []), device=self.device, dtype=torch.float32) # Find indices in the hand joints dof_idx = [] for dof_name in dof_names: if dof_name in joint_names: dof_idx.append(joint_names.index(dof_name)) else: logger.info(f"Warning: DOF {dof_name} not found in robot joints") # Support both "mode" (new) and "discrete" (legacy) config keys mode = prim_cfg.get("mode", None) if mode is None: # Legacy support: discrete=True → mode="discrete", discrete=False → mode="linear" mode = "discrete" if prim_cfg.get("discrete", False) else "linear" self._finger_primitive_map.append( { "action_name": action_name, "dof_names": dof_names, "dof_idx": dof_idx, "pos_0": pos_0, "pos_1": pos_1, "mode": mode, } ) logger.info(f"Finger primitive '{action_name}': {len(dof_names)} DOFs, mode={mode}") self._num_finger_primitives = len(self._finger_primitive_map) logger.info(f"Initialized {self._num_finger_primitives} finger primitives") def _convert_primitive_to_finger_actions(self, primitive_actions: torch.Tensor) -> torch.Tensor: """Convert primitive actions (num_envs, num_primitives) to finger joint targets (num_envs, 14). Follows groot_backup implementation: - "linear" mode: maps [-1, 1] → [0, 1] via (x + 1) / 2, then lerps between pos_0 and pos_1 - "discrete" mode: action >= 0 → pos_1 (closed), action < 0 → pos_0 (open) Args: primitive_actions: Tensor of shape (num_envs, num_primitives), values in [-1, 1] Returns: Tensor of shape (num_envs, num_finger_joints) with joint position targets """ num_envs = primitive_actions.shape[0] num_finger_joints = ( len(self._hand_joint_indices) if self._hand_joint_indices is not None else 14 ) finger_targets = torch.zeros( num_envs, num_finger_joints, device=self.device, dtype=primitive_actions.dtype ) for i, prim_cfg in enumerate(self._finger_primitive_map): action = primitive_actions[:, i] # (num_envs,), values in [-1, 1] p0 = prim_cfg["pos_0"] p1 = prim_cfg["pos_1"] dof_idx = prim_cfg["dof_idx"] mode = prim_cfg.get("mode", "linear") if mode == "linear": # Clamp to [-1, 1] then map to [0, 1] action_clamped = action.clamp(min=-1.0, max=1.0) t = (action_clamped + 1.0) / 2.0 # (num_envs,) joint_targets = torch.lerp(p0.unsqueeze(0), p1.unsqueeze(0), t.unsqueeze(1)) elif mode == "discrete": # action >= 0 → closed (pos_1), action < 0 → open (pos_0) joint_targets = torch.where( action.unsqueeze(1) >= 0, p1.unsqueeze(0), p0.unsqueeze(0) ) else: raise ValueError(f"Invalid finger primitive mode: {mode}") # Map to the correct indices in finger_targets # dof_idx are absolute joint indices, need to convert to relative hand indices for j, abs_idx in enumerate(dof_idx): if self._hand_joint_indices is not None: # Convert tensor to list if needed for .index() lookup hand_indices_list = ( self._hand_joint_indices.tolist() if isinstance(self._hand_joint_indices, torch.Tensor) else self._hand_joint_indices ) rel_idx = hand_indices_list.index(abs_idx) finger_targets[:, rel_idx] = joint_targets[:, j] else: # Fallback: assume hand joints are at the end finger_targets[:, abs_idx - 29] = joint_targets[:, j] return finger_targets def _prepare_obs_for_action_transform_module(self, obs_dict): """Use policy_atm observations for ATM if DOF mismatch exists, else use policy.""" if not self._use_policy_atm_group or "policy_atm" not in obs_dict: atm_obs_dict = obs_dict.copy() else: atm_obs_dict = obs_dict.copy() atm_obs_dict["actor_obs"] = atm_obs_dict["policy_atm"] # Ensure all observations have a sequence dimension [num_envs, seq_len, dim] # The action_transform_module expects 3D tensors for k, v in atm_obs_dict.items(): if isinstance(v, torch.Tensor) and v.dim() == 2: atm_obs_dict[k] = v.unsqueeze(1) # Add seq_len=1 dimension return atm_obs_dict def reset_all(self, global_rank=0): # noqa: ARG002 return self.reset() def process_raw_obs(self, obs, flatten_dict_obs): new_obs = { "actor_obs": obs["policy"], "critic_obs": obs["critic"], } for k, v in obs.items(): if k not in ["policy", "critic"]: if isinstance(v, dict) and flatten_dict_obs: if k == "height_map": # Special case: do not flatten height map new_obs[k] = v["height_map"] continue if k == "camera_rgb": # Special case: do not flatten camera RGB image # Keep original shape [B, H, W, C] for vision encoder new_obs[k] = v["camera_rgb"] continue obs_names = self.env.observation_manager._group_obs_term_names[k] # noqa: SLF001 new_obs[k] = torch.cat( [v[obs_name].reshape(v[obs_name].shape[0], -1) for obs_name in obs_names], dim=-1, ) else: new_obs[k] = v return new_obs def reset(self, flatten_dict_obs=True): obs, info = self.env.reset() new_obs = self.process_raw_obs(obs, flatten_dict_obs) # Initialize success_lift to False for all envs after reset (used unconditionally in step()) self.env.success_lift = torch.zeros( self.env.num_envs, dtype=torch.bool, device=self.env.device ) if self.action_transform_module is not None: # Store obs for action_transform_module when obs_dict is not provided in step() self._last_obs_dict = new_obs # Initialize last meta action buffer (policy output: latent + primitives) # meta_action_dim = tokenizer_action_dim + hand_action_dim (e.g., 64 + 2 = 66) meta_action_dim = self.config.get("meta_action_dim", 66) self.env._last_meta_action = torch.zeros( # noqa: SLF001 self.env.num_envs, meta_action_dim, dtype=torch.float32, device=self.env.device ) # Previous meta action buffer for meta_action_rate_l2 reward (token smoothness) self.env._prev_meta_action = torch.zeros( # noqa: SLF001 self.env.num_envs, meta_action_dim, dtype=torch.float32, device=self.env.device ) # Full latent buffers for full_latent_rate_l2 reward (decoder input smoothness) # tokenizer_action_dim = latent_dim (e.g., 64 = num_tokens * token_dim) tokenizer_action_dim = self.config.get("tokenizer_action_dim", 64) self.env._full_latent = torch.zeros( # noqa: SLF001 self.env.num_envs, tokenizer_action_dim, dtype=torch.float32, device=self.env.device ) self.env._prev_full_latent = torch.zeros( # noqa: SLF001 self.env.num_envs, tokenizer_action_dim, dtype=torch.float32, device=self.env.device ) # Apply camera extrinsics randomization on every reset self.apply_random_camera_extrinsics() return new_obs def apply_random_camera_extrinsics(self): """Apply per-environment random camera extrinsics (position and rotation offsets). Reads randomization ranges from config: - cameras.camera_extrinsics_randomization: Enable/disable switch (default: False) - cameras.camera_pos_rand_range: ±meters for x, y, z position - cameras.camera_roll_rand_range: ±radians for roll - cameras.camera_pitch_rand_range: ±radians for pitch - cameras.camera_yaw_rand_range: ±radians for yaw """ cameras_config = self.config.get("cameras", {}) # Check if randomization is enabled via switch if not cameras_config.get("camera_extrinsics_randomization", False): return # Get randomization ranges pos_range = cameras_config.get("camera_pos_rand_range", 0.0) roll_range = cameras_config.get("camera_roll_rand_range", 0.0) pitch_range = cameras_config.get("camera_pitch_rand_range", 0.0) yaw_range = cameras_config.get("camera_yaw_rand_range", 0.0) # Check if any randomization values are non-zero if pos_range == 0 and roll_range == 0 and pitch_range == 0 and yaw_range == 0: return # Get camera attached link from config camera_attached_link = cameras_config.get("camera_attached_link", None) if camera_attached_link is None: logger.info("Skipping camera extrinsics randomization: no camera_attached_link configured") return # Only print details on first call if not self._camera_extrinsics_randomized: logger.info("Applying random camera extrinsics per environment:") logger.info(f" Position range: ±{pos_range*100:.1f}cm") logger.info(f" Roll range: ±{roll_range:.3f} rad ({np.degrees(roll_range):.1f}°)") logger.info(f" Pitch range: ±{pitch_range:.3f} rad ({np.degrees(pitch_range):.1f}°)") logger.info(f" Yaw range: ±{yaw_range:.3f} rad ({np.degrees(yaw_range):.1f}°)") stage = omni.usd.get_context().get_stage() for env_id in range(self.env.scene.num_envs): camera_prim_path = f"/World/envs/env_{env_id}/Robot/{camera_attached_link}/ego_camera" camera_prim = stage.GetPrimAtPath(camera_prim_path) if not camera_prim.IsValid(): if env_id == 0: logger.info(f" Warning: Camera prim not found at {camera_prim_path}") continue # Get current camera transform xformable = UsdGeom.Xformable(camera_prim) # On first call, store the original/base transforms if env_id not in self._camera_base_transforms: base_translate = None base_orient = None for op in xformable.GetOrderedXformOps(): if op.GetOpType() == UsdGeom.XformOp.TypeTranslate: base_translate = Gf.Vec3d(op.Get()) # Make a copy elif op.GetOpType() == UsdGeom.XformOp.TypeOrient: base_orient = Gf.Quatd(op.Get()) # Make a copy self._camera_base_transforms[env_id] = { "translate": base_translate, "orient": base_orient, } # Get base transforms base_translate = self._camera_base_transforms[env_id]["translate"] base_orient = self._camera_base_transforms[env_id]["orient"] # Sample random deltas if pos_range > 0: # noqa: SIM108 pos_delta = np.random.uniform(-pos_range, pos_range, 3) # noqa: NPY002 else: pos_delta = np.zeros(3) # Sample rotation deltas (roll, pitch, yaw) roll_delta = np.random.uniform(-roll_range, roll_range) if roll_range > 0 else 0.0 # noqa: NPY002 pitch_delta = np.random.uniform(-pitch_range, pitch_range) if pitch_range > 0 else 0.0 # noqa: NPY002 yaw_delta = np.random.uniform(-yaw_range, yaw_range) if yaw_range > 0 else 0.0 # noqa: NPY002 # Apply position delta relative to BASE (not current) if base_translate is not None: new_pos = Gf.Vec3d( base_translate[0] + pos_delta[0], base_translate[1] + pos_delta[1], base_translate[2] + pos_delta[2], ) # Find and update the translate op for op in xformable.GetOrderedXformOps(): if op.GetOpType() == UsdGeom.XformOp.TypeTranslate: op.Set(new_pos) break # Apply rotation delta relative to BASE (not current) if base_orient is not None and (roll_delta != 0 or pitch_delta != 0 or yaw_delta != 0): # Convert euler deltas to quaternion # Order: roll (X), pitch (Y), yaw (Z) cr, sr = np.cos(roll_delta / 2), np.sin(roll_delta / 2) cp, sp = np.cos(pitch_delta / 2), np.sin(pitch_delta / 2) cy, sy = np.cos(yaw_delta / 2), np.sin(yaw_delta / 2) # Quaternion from euler (ZYX convention) delta_quat = Gf.Quatd( cr * cp * cy + sr * sp * sy, # w sr * cp * cy - cr * sp * sy, # x cr * sp * cy + sr * cp * sy, # y cr * cp * sy - sr * sp * cy, # z ) # Compose: new = delta * BASE (not current!) new_orient = delta_quat * base_orient # Find and update the orient op for op in xformable.GetOrderedXformOps(): if op.GetOpType() == UsdGeom.XformOp.TypeOrient: op.Set(new_orient) break # Store the random deltas for each environment (for projection in observations) if not hasattr(self, "_camera_random_deltas"): self._camera_random_deltas = {} self._camera_random_deltas[env_id] = { "pos_delta": pos_delta.tolist() if isinstance(pos_delta, np.ndarray) else [0, 0, 0], "roll_delta": float(roll_delta), "pitch_delta": float(pitch_delta), "yaw_delta": float(yaw_delta), } if env_id == 0 and not self._camera_extrinsics_randomized: logger.info( f" Env 0: pos_delta={pos_delta}, rot_delta=[{roll_delta:.3f}, {pitch_delta:.3f}, {yaw_delta:.3f}]" # noqa: E501 ) # Mark as initialized (for print suppression) self._camera_extrinsics_randomized = True def _decode_direct_latent(self, full_latent, atm_obs_dict): """Decode full latent directly using ATM decoder (skip encoder). Used for student rollout where policy outputs full latent. Args: full_latent: (batch, latent_dim) full latent from student policy atm_obs_dict: observation dict for ATM Returns: body_actions: (batch, seq_len, action_dim) decoded actions """ # noqa: D205 # Get proprioception for decoder if "policy_atm" in atm_obs_dict: # noqa: SIM108 proprioception = atm_obs_dict["policy_atm"] else: proprioception = atm_obs_dict["actor_obs"] # Ensure sequence dimension if proprioception.dim() == 2: proprioception = proprioception.unsqueeze(1) atm = self.action_transform_module.actor_module return self._decode_direct_latent_batch(full_latent, proprioception, atm) def _decode_direct_latent_batch(self, full_latent, proprioception, atm): """Decode a batch of full latents directly using ATM decoder. The student policy outputs pre-quantization values (latent + residual). We apply quantization here before decoding to match the teacher's flow: - Teacher: latent + residual -> quantize -> decode - Student: direct_latent -> quantize -> decode Args: full_latent: (batch, latent_dim) full latent (pre-quantization) proprioception: (batch, seq, proprio_dim) proprioception input atm: ATM actor module Returns: body_actions: (batch, seq_len, action_dim) decoded actions """ batch_size = full_latent.shape[0] # Reshape latent for quantization: (batch, latent_dim) -> (batch, num_tokens, token_dim) latent_reshaped = full_latent.view(batch_size, atm.max_num_tokens, atm.token_dim) # Apply quantization (same as teacher's flow) # This ensures student inference matches teacher: latent -> quantize -> decode if atm.quantizer is not None: quantized_codes, _ = atm.quantizer(latent_reshaped) tokens_for_decode = quantized_codes.contiguous() else: tokens_for_decode = latent_reshaped # Reshape for decoder: (batch, num_tokens, token_dim) -> (batch, 1, num_tokens, token_dim) tokens_reshaped = tokens_for_decode.unsqueeze(1) tokens_flattened = tokens_for_decode.view(batch_size, -1).unsqueeze(1) # Prepare decode input decode_input_dict = { "token": tokens_reshaped, "token_flattened": tokens_flattened, # (batch, 1, latent_dim) "proprioception": proprioception, } # Decode directly (skip encoding entirely) decoded_output = atm.decode("g1_dyn", decode_input_dict) # Get body actions from decoder output # Note: "meta_action" is for special hierarchical policies, "action" is standard body_actions = decoded_output.get("meta_action", decoded_output.get("action")) if body_actions is None: raise KeyError( f"Decoder output missing 'action' or 'meta_action'. Keys: {decoded_output.keys()}" ) return body_actions def step(self, actions): if self.action_transform_module is not None: # Use provided obs_dict or fall back to stored obs from last reset/step if "obs_dict" in actions: obs_dict = actions["obs_dict"].copy() else: # Fallback for callbacks that don't provide obs_dict (e.g., im_eval) obs_dict = getattr(self, "_last_obs_dict", None) if obs_dict is None: raise ValueError( "action_transform_module requires obs_dict but none was provided or stored" ) obs_dict = obs_dict.copy() if isinstance(obs_dict, dict) else obs_dict meta_actions = actions["actions"] # Determine action mode: "direct_latent", "residual", or "mixed" # Priority: 1) explicit action_mode in actions dict, 2) config flag # During training: trainer sets action_mode explicitly # During eval: fallback to config flags action_mode = actions.get("action_mode", None) if action_mode is None: # Fallback for eval scripts that don't set action_mode if self._use_student_direct_latent: action_mode = "direct_latent" elif self._use_latent_residual: action_mode = "residual" else: action_mode = "residual" # Default to residual if nothing specified # Shift meta action buffers for meta_action_rate_l2 reward (token smoothness) self.env._prev_meta_action = self.env._last_meta_action.clone() # noqa: SLF001 # Store meta action for observation (last policy output) self.env._last_meta_action = meta_actions.clone() # noqa: SLF001 atm_obs_dict = self._prepare_obs_for_action_transform_module(obs_dict) # Split actions: first tokenizer_action_dim for tokenizer, rest for hands tokenizer_action_dim = self.config.get("tokenizer_action_dim") tokenizer_meta_actions = meta_actions[:, :tokenizer_action_dim] hand_actions_raw = meta_actions[:, tokenizer_action_dim:] # Override hand actions with motion data if configured if self.config.get("use_motion_hand_actions", False): motion_cmd = self.env.command_manager.get_term("motion") left_action = motion_cmd.get_hand_action("left_hand") right_action = motion_cmd.get_hand_action("right_hand") if left_action is None or right_action is None: raise ValueError( "use_motion_hand_actions=True but hand_action_left/right not found in motion data. " "Ensure processed_robot_motions.pkl contains 'hand_action_left' and 'hand_action_right' arrays." # noqa: E501 ) # Use motion data directly: -1.0 = open, +1.0 = closed # Threshold at 0 in _convert_primitive_to_finger_actions hand_actions_raw = torch.stack([left_action, right_action], dim=-1) # Convert primitive actions to finger joint targets if enabled if self._use_finger_primitive and self._finger_primitive_map: # Store raw primitive actions on env for reward computation (before clamping) self.env._finger_primitive_actions_raw = hand_actions_raw # noqa: SLF001 hand_actions = self._convert_primitive_to_finger_actions(hand_actions_raw) else: self.env._finger_primitive_actions_raw = None # noqa: SLF001 hand_actions = hand_actions_raw if action_mode == "direct_latent": # Student direct latent mode: policy outputs FULL latent, not residual # Skip ATM encoder entirely - go directly to decoder body_actions = self._decode_direct_latent(tokenizer_meta_actions, atm_obs_dict) elif action_mode == "residual": # Teacher/residual mode: policy outputs residual that's added to ATM encoded tokens # Apply scaling to residual before passing to ATM scaled_residual = tokenizer_meta_actions * self._latent_residual_scale # Add residual in latent/token space (after encoding, before decoding) body_actions = self.action_transform_module( atm_obs_dict, latent_residual=scaled_residual, latent_residual_mode=self._latent_residual_mode, ) elif action_mode == "mixed": # Mixed rollout: some envs use teacher (residual), some use student (direct_latent) # is_teacher_env is a boolean mask: True = teacher/residual, False = student/direct_latent is_teacher_env = actions.get("is_teacher_env") if is_teacher_env is None: raise ValueError( "action_mode='mixed' requires 'is_teacher_env' mask in actions dict" ) num_envs = tokenizer_meta_actions.shape[0] atm = self.action_transform_module.actor_module # Get proprioception for decoder (needed for student/direct_latent mode) if "policy_atm" in atm_obs_dict: proprioception = atm_obs_dict["policy_atm"] else: proprioception = atm_obs_dict["actor_obs"] if proprioception.dim() == 2: proprioception = proprioception.unsqueeze(1) teacher_mask = is_teacher_env student_mask = ~is_teacher_env num_teacher = teacher_mask.sum().item() num_student = student_mask.sum().item() # Initialize placeholders teacher_indices = None student_indices = None teacher_body_actions = None student_body_actions = None # Process teacher envs (residual mode) if any if num_teacher > 0: teacher_indices = teacher_mask.nonzero(as_tuple=True)[0] teacher_latent = tokenizer_meta_actions[teacher_indices] # Prepare teacher obs dict (subset of envs) # Only include keys that ATM actually needs: tokenizer and actor_obs teacher_atm_obs = {} atm_keys = ["tokenizer", "actor_obs"] for obs_key in atm_keys: if obs_key in atm_obs_dict: teacher_atm_obs[obs_key] = atm_obs_dict[obs_key][teacher_indices] scaled_residual = teacher_latent * self._latent_residual_scale teacher_body_actions = self.action_transform_module( teacher_atm_obs, latent_residual=scaled_residual, latent_residual_mode=self._latent_residual_mode, ) # Process student envs (direct_latent mode) if any if num_student > 0: student_indices = student_mask.nonzero(as_tuple=True)[0] student_latent = tokenizer_meta_actions[student_indices] student_proprio = proprioception[student_indices] student_body_actions = self._decode_direct_latent_batch( student_latent, student_proprio, atm ) # Merge results - determine output shape from whichever mode ran if teacher_body_actions is not None: out_seq_len = teacher_body_actions.shape[1] out_action_dim = teacher_body_actions.shape[2] dtype = teacher_body_actions.dtype elif student_body_actions is not None: out_seq_len = student_body_actions.shape[1] out_action_dim = student_body_actions.shape[2] dtype = student_body_actions.dtype else: raise RuntimeError("Mixed mode: no envs to process (both masks empty)") body_actions = torch.zeros( num_envs, out_seq_len, out_action_dim, device=self.device, dtype=dtype ) if teacher_body_actions is not None and teacher_indices is not None: body_actions[teacher_indices] = teacher_body_actions if student_body_actions is not None and student_indices is not None: body_actions[student_indices] = student_body_actions else: raise ValueError( f"Unknown action_mode: {action_mode}. " f"Valid modes are 'direct_latent', 'residual', or 'mixed'." ) # Store full latent (decoder input) for full_latent_rate_l2 reward # Only needed for residual mode (teacher RL training); student modes # (direct_latent, mixed) use L2 distillation loss, not RL rewards. if action_mode == "residual": self.env._prev_full_latent = self.env._full_latent.clone() # noqa: SLF001 atm_module = self.action_transform_module.actor_module if ( hasattr(atm_module, "_last_full_latent_flat") and atm_module._last_full_latent_flat is not None # noqa: SLF001 ): fl = atm_module._last_full_latent_flat # noqa: SLF001 if fl.dim() == 3: fl = fl[:, -1, :] # (batch, latent_dim) self.env._full_latent = fl.to(self.env.device) # noqa: SLF001 body_actions = body_actions[:, -1] # Take last timestep if ( self._body_joint_indices is not None and self._hand_joint_indices is not None and len(self._body_joint_indices) > 0 and len(self._hand_joint_indices) > 0 ): num_envs = body_actions.shape[0] env_actions = torch.zeros( num_envs, self._current_num_joints, device=self.device, dtype=body_actions.dtype ) env_actions[:, self._body_joint_indices] = body_actions env_actions[:, self._hand_joint_indices] = hand_actions else: env_actions = torch.cat([body_actions, hand_actions], dim=-1) else: env_actions = actions["actions"] action_clip_value = self.config.get("action_clip_value", None) if action_clip_value is not None and action_clip_value > 0: env_actions = torch.clip(env_actions, -action_clip_value, action_clip_value) # Lightweight action plot update (env 0, first N joints) if self.turn_on_visualization: try: if self._action_fig is None: self._init_action_plot(env_actions.shape[-1]) if (self._step_counter % self._viz_every_n_steps) == 0: self._update_action_plot(env_actions) self._step_counter += 1 except Exception: # noqa: S110, BLE001 pass obs_dict, rew, terminated, truncated, extras = self.env.step(env_actions) # compute dones for compatibility with RSL-RL dones = (terminated | truncated).to(dtype=torch.long) # Zero out action/latent rate buffers for envs that just reset # This prevents a false large rate penalty on the first step of a new episode # Only applies when action_transform_module is used (buffers created in reset()) reset_mask = dones.bool() if reset_mask.any() and hasattr(self.env, "_prev_meta_action"): self.env._prev_meta_action[reset_mask] = 0.0 # noqa: SLF001 self.env._last_meta_action[reset_mask] = 0.0 # noqa: SLF001 self.env._prev_full_latent[reset_mask] = 0.0 # noqa: SLF001 self.env._full_latent[reset_mask] = 0.0 # noqa: SLF001 # Compute success_lift metric: check if object has no contact with table (lifted) # Skip frames before first contact (from contact label data) if ( hasattr(self.env, "scene") and "object_to_table_contact_sensor" in self.env.scene.sensors ): from isaaclab.sensors import ContactSensor sensor: ContactSensor = self.env.scene["object_to_table_contact_sensor"] contact_force = sensor.data.force_matrix_w # [num_envs, 1, 1, 3] force_magnitude = torch.norm(contact_force, dim=-1).sum(dim=(-1, -2)) # [num_envs] contact_force_threshold = self.config.get("lift_contact_force_threshold", 0.1) if hasattr(self.env, "success_lift"): self.env.success_lift = self.env.success_lift & (~dones.bool()) # Get first contact frame from motion command (contact label data) is_before_contact = torch.ones( self.env.num_envs, dtype=torch.bool, device=self.env.device ) if self.motion_command is not None: per_env_first_contact = getattr(self.motion_command, "_per_env_first_contact", None) if per_env_first_contact is not None: current_time = ( self.motion_command.motion_start_time_steps + self.motion_command.time_steps ) is_before_contact = current_time < per_env_first_contact # Object is currently lifted if there's no contact force from table is_currently_lifted = force_magnitude <= contact_force_threshold # Only update success_lift after first contact (cumulative OR) # Before first contact, success_lift stays False regardless of contact self.env.success_lift = torch.where( is_before_contact, self.env.success_lift, # Keep current value (False after reset) is_currently_lifted | self.env.success_lift, ) else: # If no contact sensor, set to False for all envs self.env.success_lift = torch.zeros( self.env.num_envs, dtype=torch.bool, device=self.env.device ) extras["time_outs"] = truncated extras["episode"] = {} extras["to_log"] = {} for k, v in extras["log"].items(): if isinstance(v, torch.Tensor): extras["to_log"][k] = v else: extras["to_log"][k] = torch.tensor(v, dtype=torch.float) if self._motion_lib is not None and self._motion_lib.use_adaptive_sampling: extras["to_log"][ "adp_samp/num_episodes_min" ] = self._motion_lib.adp_samp_num_episodes.min() extras["to_log"][ "adp_samp/num_episodes_max" ] = self._motion_lib.adp_samp_num_episodes.max() extras["to_log"][ "adp_samp/num_episodes_mean" ] = self._motion_lib.adp_samp_num_episodes.mean() extras["to_log"][ "adp_samp/num_failures_min" ] = self._motion_lib.adp_samp_num_failures.min() extras["to_log"][ "adp_samp/num_failures_max" ] = self._motion_lib.adp_samp_num_failures.max() extras["to_log"][ "adp_samp/num_failures_mean" ] = self._motion_lib.adp_samp_num_failures.mean() extras["to_log"][ "adp_samp/failure_rate_min" ] = self._motion_lib.adp_samp_failure_rate_raw.min() extras["to_log"][ "adp_samp/failure_rate_max" ] = self._motion_lib.adp_samp_failure_rate_raw.max() extras["to_log"][ "adp_samp/failure_rate_mean" ] = self._motion_lib.adp_samp_failure_rate_raw.mean() if hasattr(self._motion_lib, "adp_sampling_active_prob"): prob = self._motion_lib.adp_sampling_active_prob uniform_prob = 1.0 / len(prob) if len(prob) > 0 else 1.0 extras["to_log"]["adp_samp/prob_max"] = prob.max() extras["to_log"]["adp_samp/prob_min"] = prob.min() extras["to_log"]["adp_samp/prob_mean"] = prob.mean() extras["to_log"]["adp_samp/prob_max_over_uniform"] = prob.max() / uniform_prob extras["to_log"]["adp_samp/effective_num_bins"] = 1.0 / (prob**2).sum() # How many bins have prob > 10x uniform (significantly concentrated) # Note: max allowed is 50x uniform, so 10x is 20% of the cap extras["to_log"]["adp_samp/num_concentrated_bins"] = ( (prob > 10 * uniform_prob).sum().float() ) eps_mean = self._motion_lib.adp_samp_num_episodes.mean() if eps_mean > 0: extras["to_log"]["adp_samp/episodes_max_over_mean"] = ( self._motion_lib.adp_samp_num_episodes.max() / eps_mean ) new_obs = self.process_raw_obs(obs_dict, flatten_dict_obs=True) # Store obs for action_transform_module when obs_dict is not provided in next step() self._last_obs_dict = new_obs self.extras = extras # Store env_actions for callbacks (e.g., MultiLatentSaveCallback) extras["env_actions"] = env_actions.detach().cpu() return new_obs, rew, dones, extras def get_env_data(self, key): if key == "ref_body_pos_extend": return self.motion_command.robot_body_pos_w elif key == "rigid_body_pos_extend": return self.motion_command.body_pos_w else: return self.env.get_env_data(key) def render_results(self): pass def end_render_results(self): if "render_envs" in self.env.recorder_manager._terms: # noqa: SLF001 self.env.recorder_manager._terms["render_envs"].close_writers() # noqa: SLF001 if "trajectory" in self.env.recorder_manager._terms: # noqa: SLF001 self.env.recorder_manager._terms["trajectory"].close_writers() # noqa: SLF001 if self._action_fig is not None: try: import matplotlib.pyplot as plt plt.close(self._action_fig) except Exception: # noqa: S110, BLE001 pass self._action_fig = None self._action_lines = None self._action_hist = None self._hist_idx = 0 self._blit_background = None def set_is_evaluating(self, is_evaluating: bool = True, global_rank=0, **_kwargs): self.is_evaluating = is_evaluating if self.motion_command is not None: self.motion_command.set_is_evaluating(is_evaluating) if self.force_command is not None: self.force_command.set_is_evaluating(is_evaluating) if is_evaluating: self.begin_seq_motion_samples(global_rank) def begin_seq_motion_samples(self, global_rank=0): logger.info("Loading motions for evaluation") self.start_idx = global_rank * self.num_envs self._motion_lib.load_motions_for_evaluation(start_idx=self.start_idx) self.reset_all(global_rank=global_rank) def forward_motion_samples(self, global_rank=0, world_size=1): old_start_idx = self.start_idx self.start_idx += world_size * self.num_envs logger.info( f"Forward motions for evaluation from {old_start_idx} to {self.start_idx} - rank: {global_rank} - world size: {world_size}" # noqa: E501 ) self._motion_lib.load_motions_for_evaluation(start_idx=self.start_idx) self.reset_all(global_rank=global_rank) def focusing_viewer(self): if not self.viewer_focused: # Focus on robot asset (tracking mode) self.env.viewport_camera_controller.update_view_to_world() self.env.viewport_camera_controller.update_view_to_asset_root("robot") self.viewer_focused = True logger.info("Focused on robot") else: # Switch to free camera mode centered on robot self.env.viewport_camera_controller.viewer_origin = torch.zeros_like( self.env.scene["robot"].data.root_pos_w[0] ) self.env.viewport_camera_controller.cfg.origin_type = "world" cam_eye = self.env.scene["robot"].data.root_pos_w[0].cpu().numpy() + np.array( [2.0, 2.0, 1.5] ) cam_target = self.env.scene["robot"].data.root_pos_w[0].cpu().numpy() self.env.viewport_camera_controller._env.sim.set_camera_view( # noqa: SLF001 eye=cam_eye, target=cam_target ) self.viewer_focused = False logger.info("Switched to free camera mode") # noqa: RUF100, T201 def set_is_training(self, **_kwargs): self.is_evaluating = False self.resample_motion() def resample_motion(self): res = self._motion_lib.load_motions_for_training( max_num_seqs=min(self.num_envs, self.motion_command.max_num_load_motions) ) if res: self.reset_all() else: logger.info("No new motions loaded, skipping reset") # noqa: RUF100, T201 def sync_and_compute_adaptive_sampling(self, accelerator, sync_across_gpus=False): if self._motion_lib is not None: self._motion_lib.sync_and_compute_adaptive_sampling( accelerator, sync_across_gpus=sync_across_gpus ) def load_env_state_dict(self, state_dict): if "motion_lib" in state_dict: self._motion_lib.load_state_dict(state_dict["motion_lib"]) if self._motion_lib.use_adaptive_sampling: self.resample_motion() def get_env_state_dict(self): state_dict = { "motion_lib": self._motion_lib.get_state_dict(), } return state_dict def reinit_dr(self, **_kwargs): pass def setup_keyboard(self): try: from isaaclab.devices.keyboard.se2_keyboard import Se2Keyboard, Se2KeyboardCfg cfg = Se2KeyboardCfg() self.keyboard_interface = Se2Keyboard(cfg) self.keyboard_interface.add_callback("R", self.reset_all) self.keyboard_interface.add_callback("T", self.forward_motion_samples) self.keyboard_interface.add_callback("F", self.focusing_viewer) self.keyboard_interface.add_callback("V", self.toggle_debug_vis) except Exception as e: # noqa: BLE001 logger.info(f"Error setting up keyboard: {e}") # noqa: RUF100, T201 def toggle_debug_vis(self): if self.motion_command is not None and hasattr(self.motion_command, "_set_debug_vis_impl"): self._debug_vis_enabled = not getattr(self, "_debug_vis_enabled", True) self.motion_command._set_debug_vis_impl(self._debug_vis_enabled) # noqa: SLF001 logger.info(f"Debug visualization: {'ON' if self._debug_vis_enabled else 'OFF'}") # noqa: RUF100, T201 # --- Action plotting helpers --- def _init_action_plot(self, action_dim: int): plot_dim = min(int(self._plot_action_dim), int(action_dim)) if plot_dim <= 0: return try: import matplotlib.pyplot as plt except Exception as e: # noqa: BLE001, F841 return plt.ion() import math cols = math.ceil(math.sqrt(plot_dim)) rows = math.ceil(plot_dim / cols) fig, axes = plt.subplots( rows, cols, sharex=True, sharey=True, figsize=(cols * 3.0, rows * 2.2) ) axes = np.array(axes).reshape(-1) self._plot_window = max(10, int(self._plot_window)) x_vals = np.arange(self._plot_window) self._action_hist = np.zeros((self._plot_window, plot_dim), dtype=np.float32) lines = [] y_min, y_max = -float(self._action_ylim), float(self._action_ylim) for i in range(rows * cols): ax_i = axes[i] if i < plot_dim: (ln,) = ax_i.plot(x_vals, self._action_hist[:, i], linewidth=1.0) ln.set_animated(True) lines.append(ln) ax_i.set_ylim((y_min, y_max)) ax_i.set_xlim((0, self._plot_window - 1)) ax_i.set_title(f"J{i}", fontsize=8) if i // cols == rows - 1: ax_i.set_xlabel("step", fontsize=8) if i % cols == 0: ax_i.set_ylabel("act", fontsize=8) else: ax_i.axis("off") fig.suptitle(f"Action time series (env 0) - first {plot_dim} joints", fontsize=10) fig.tight_layout(rect=(0, 0.02, 1, 0.96)) self._action_fig = fig self._action_lines = lines self._hist_idx = -1 try: fig.canvas.draw() self._blit_background = fig.canvas.copy_from_bbox(fig.bbox) except Exception: # noqa: BLE001 self._blit_background = None self._blit_supported = False fig.canvas.flush_events() def _update_action_plot(self, env_actions: torch.Tensor): if self._action_lines is None or self._action_hist is None: return plot_dim = len(self._action_lines) with torch.no_grad(): vals = env_actions[0, :plot_dim].detach().to("cpu").numpy().astype(np.float32) vals = np.clip(vals, -self._action_ylim, self._action_ylim) # advance circular buffer self._hist_idx = (self._hist_idx + 1) % self._plot_window self._action_hist[self._hist_idx, :plot_dim] = vals # display in chronological order using a view with roll disp = np.roll(self._action_hist, shift=-(self._hist_idx + 1), axis=0) for i, ln in enumerate(self._action_lines): ln.set_ydata(disp[:, i]) if self._action_fig is not None: canvas = self._action_fig.canvas # Periodically refresh background to handle resizes or overdraw if self._blit_supported and ( self._blit_background is None or (self._step_counter % max(1, self._blit_refresh_interval) == 0) ): try: self._action_fig.canvas.draw() self._blit_background = canvas.copy_from_bbox(self._action_fig.bbox) except Exception: # noqa: BLE001 self._blit_background = None self._blit_supported = False if self._blit_supported and self._blit_background is not None: try: canvas.restore_region(self._blit_background) for ln in self._action_lines: ln.axes.draw_artist(ln) canvas.blit(self._action_fig.bbox) canvas.flush_events() return except Exception: # noqa: BLE001 self._blit_supported = False # Fallback full redraw canvas.draw_idle() canvas.flush_events() @property def motion_ids(self): return self.motion_command.motion_ids def setup_replay_grid(self, spacing=2.0, rows=None, cols=None): """Setup a custom grid layout for environment origins during replay. Args: spacing: Distance between environments in meters (default: 2.0) rows: Number of rows in the grid (default: auto-calculated) cols: Number of columns in the grid (default: auto-calculated) Returns: Tensor of shape (num_envs, 3) with custom origins """ import math num_envs = self.num_envs # Auto-calculate grid dimensions if not provided if rows is None and cols is None: # Try to make a square-ish grid cols = int(math.ceil(math.sqrt(num_envs))) # noqa: RUF046 rows = int(math.ceil(num_envs / cols)) # noqa: RUF046 elif rows is None: rows = int(math.ceil(num_envs / cols)) # noqa: RUF046 elif cols is None: cols = int(math.ceil(num_envs / rows)) # noqa: RUF046 logger.info(f"Setting up replay grid: {rows} rows x {cols} cols with {spacing}m spacing") # Create grid origins custom_origins = torch.zeros((num_envs, 3), device=self.device, dtype=torch.float32) for i in range(num_envs): row = i // cols col = i % cols # Center the grid around origin x_offset = (col - (cols - 1) / 2.0) * spacing y_offset = (row - (rows - 1) / 2.0) * spacing custom_origins[i, 0] = x_offset custom_origins[i, 1] = y_offset custom_origins[i, 2] = 0.0 # Keep z at ground level # Store custom origins self._replay_custom_origins = custom_origins logger.info( f"Grid bounds: X=[{custom_origins[:, 0].min():.1f}, {custom_origins[:, 0].max():.1f}], " f"Y=[{custom_origins[:, 1].min():.1f}, {custom_origins[:, 1].max():.1f}]" ) return custom_origins def run_replay( self, motion_id=None, start_time_step=0, speed=1.0, loop=True, enable_vis=False, target_fps=50, grid_spacing=2.0, grid_rows=None, grid_cols=None, save_video_path=None, ): """Run a complete motion replay with automatic rendering loop for all environments. This is the high-level convenience function that handles everything. Args: motion_id: The motion ID(s) to replay. Can be: - None: uses current motion_ids for all envs - int: single motion ID for all envs - list/tensor: motion ID per environment start_time_step: Starting time step in the motion (default: 0) speed: Playback speed multiplier (default: 1.0) loop: Whether to loop the motion (default: True) enable_vis: Whether to enable debug visualization (default: False) target_fps: Target frames per second for rendering (default: 50) grid_spacing: Distance between environments in meters (default: 2.0) grid_rows: Number of rows in grid layout (default: auto) grid_cols: Number of columns in grid layout (default: auto) save_video_path: Path to save video file (default: None, no video saved) Requires render_results=True in config to enable eval_camera. Returns: List of dictionaries with motion metadata for each environment """ # noqa: D205 import os import time # Setup custom grid origins self.setup_replay_grid(spacing=grid_spacing, rows=grid_rows, cols=grid_cols) # Initialize replay info = self.setup_replay_motion( motion_id=motion_id, start_time_step=start_time_step, speed=speed, loop=loop, enable_vis=enable_vis, ) if info is None: return None # Initialize video writer if save_video_path is specified video_writer = None if save_video_path is not None: if "eval_camera" not in self.env.scene.sensors: logger.info("WARNING: save_video_path specified but eval_camera not available.") logger.info(" Set ++manager_env.config.render_results=True to enable it.") else: import imageio os.makedirs(os.path.dirname(os.path.abspath(save_video_path)), exist_ok=True) video_writer = imageio.get_writer( save_video_path, fps=target_fps, codec="libx264", quality=5, pixelformat="yuv420p", ) logger.info(f"[Video] Recording to: {save_video_path}") # Run the rendering loop frame_time = 1.0 / target_fps last_time = time.time() frame_count = 0 try: while self.step_replay(): # Position camera BEFORE rendering (if recording) if video_writer is not None and "eval_camera" in self.env.scene.sensors: root_pos = self.motion_command.robot.data.body_pos_w[:, 0] camera_offset = self.config.get("eval_camera_offset", [-2, -2, 1]) eye = root_pos + torch.tensor(camera_offset, device=self.device) self.env.scene["eval_camera"].set_world_poses_from_view(eye, root_pos) # Render the scene (triggers camera render in headless mode) if hasattr(self.env, "sim"): try: self.env.sim.render() except Exception as e: # noqa: BLE001 logger.info(f"Render error: {e}") # Capture video frame if recording if video_writer is not None and "eval_camera" in self.env.scene.sensors: # Update camera to refresh its data after render self.env.scene["eval_camera"].update(dt=0.0) # Grab frame from camera rgb_frame = self.env.scene["eval_camera"].data.output["rgb"] frame = rgb_frame[0].cpu().numpy().astype(np.uint8) video_writer.append_data(frame) frame_count += 1 # Frame rate limiting (only when not headless or not saving video) if not self.config.get("headless", False) or video_writer is None: current_time = time.time() elapsed = current_time - last_time sleep_time = max(0, (frame_time / speed) - elapsed) if sleep_time > 0: time.sleep(sleep_time) last_time = time.time() except KeyboardInterrupt: logger.info("\nReplay interrupted by user") # Close video writer if video_writer is not None: video_writer.close() logger.info(f"[Video] Saved {frame_count} frames to: {save_video_path}") logger.info("\nReplay complete!") return info def setup_replay_motion( self, motion_id=None, start_time_step=0, speed=1.0, loop=True, enable_vis=False ): """Initialize motion replay for all environments (low-level function for manual control). Use run_replay() instead if you want automatic rendering loop. After calling this, you need to manually call step_replay() in a loop and handle rendering yourself. For automatic rendering, use run_replay(). Args: motion_id: The motion ID(s) to replay. Can be: - None: uses current motion_ids for all envs - int: single motion ID for all envs - list/tensor: motion ID per environment start_time_step: Starting time step in the motion (default: 0) speed: Playback speed multiplier (default: 1.0) loop: Whether to loop the motion (default: True) enable_vis: Whether to enable debug visualization (default: False) Returns: List of dictionaries with motion metadata for each environment """ # noqa: D205 if self._motion_lib is None: logger.info("No motion library available for replay") return None # Replay all environments num_replay_envs = self.num_envs replay_env_ids = torch.arange(num_replay_envs, device=self.device, dtype=torch.long) # Get motion IDs if motion_id is None: motion_ids = self.motion_command.motion_ids elif isinstance(motion_id, int): motion_ids = torch.full( (num_replay_envs,), motion_id, device=self.device, dtype=torch.long ) elif isinstance(motion_id, list | tuple): if len(motion_id) != num_replay_envs: logger.info( f"Error: motion_id list length ({len(motion_id)}) doesn't match number of environments ({num_replay_envs})" # noqa: E501 ) return None motion_ids = torch.tensor(motion_id, device=self.device, dtype=torch.long) elif isinstance(motion_id, torch.Tensor): if motion_id.shape[0] != num_replay_envs: logger.info( f"Error: motion_id tensor size ({motion_id.shape[0]}) doesn't match number of environments ({num_replay_envs})" # noqa: E501 ) return None motion_ids = motion_id.to(self.device) else: logger.info(f"Error: Invalid motion_id type: {type(motion_id)}") return None # Get motion info for display num_steps_per_env = self._motion_lib.get_motion_num_steps(motion_ids) max_num_steps = num_steps_per_env.max().item() unique_motions = torch.unique(motion_ids) logger.info(f"\n{'='*60}") logger.info(f"Batch Replaying {num_replay_envs} Environments") logger.info(f"Unique motion IDs: {unique_motions.tolist()}") logger.info(f"Max frames: {max_num_steps}") logger.info(f"Max duration: {max_num_steps / 50.0:.2f}s (@ 50 FPS)") logger.info(f"Speed: {speed}x") logger.info(f"{'='*60}\n") # Enable visualization if requested if enable_vis and hasattr(self.motion_command, "_set_debug_vis_impl"): self.motion_command._set_debug_vis_impl(True) # noqa: SLF001 # Initialize contact point visualization if available and enabled self._replay_contact_visualizer = None self._replay_vis_enabled = enable_vis if enable_vis and VISUALIZATION_AVAILABLE: self._setup_contact_center_visualizer() # Replay state self._replay_active = True self._replay_paused = False self._replay_reverse = False self._replay_time_steps = torch.full( (num_replay_envs,), start_time_step, device=self.device, dtype=torch.long ) self._replay_motion_ids = motion_ids self._replay_env_ids = replay_env_ids self._replay_speed = speed self._replay_loop = loop self._replay_num_steps_per_env = num_steps_per_env self._replay_max_num_steps = max_num_steps # Pre-load table metadata from pkl files for ALL motions (per-env) # This supports multi-motion replay where each env can have different table positions self._replay_table_pos = None # Shape: (num_envs, 3) self._replay_table_quat = None # Shape: (num_envs, 4) if hasattr(self.env, "scene") and "table" in self.env.scene.rigid_objects: try: import os import joblib # Derive meta directory from motion_file path in config motion_lib_cfg = getattr(self.motion_command.cfg, "motion_lib_cfg", None) motion_file = motion_lib_cfg.get("motion_file", "") if motion_lib_cfg else "" if motion_file and "/robot" in motion_file: if os.path.isdir(motion_file): meta_dir = motion_file.replace("/robot", "/meta") else: meta_dir = os.path.dirname(motion_file).replace("/robot", "/meta") else: meta_dir = "data/motion_lib_grab/meta" # Pre-load motion file data if it's a single file (for table fallback) motion_file_data = None if motion_file and os.path.isfile(motion_file) and not os.path.isdir(meta_dir): motion_file_data = joblib.load(motion_file) # Load table metadata for each environment's assigned motion table_pos_list = [] table_quat_list = [] table_scale_list = [] # Also load scales for offset calculation loaded_count = 0 loaded_from_motion_file = 0 for env_idx in range(num_replay_envs): motion_idx = motion_ids[env_idx].item() motion_key = self._motion_lib.curr_motion_keys[motion_idx] meta_file = os.path.join(meta_dir, f"{motion_key}.pkl") if os.path.exists(meta_file): meta = joblib.load(meta_file) pos = torch.tensor( meta.get("table_pos", [0.0, 0.0, 0.8]), device=self.device ).float() quat = torch.tensor( meta.get("table_quat", [1.0, 0.0, 0.0, 0.0]), device=self.device ).float() scale = torch.tensor( meta.get("table_scale", [1.0, 1.0, 1.0]), device=self.device ).float() table_pos_list.append(pos) table_quat_list.append(quat) table_scale_list.append(scale) loaded_count += 1 elif motion_file_data is not None and motion_key in motion_file_data: # Fallback: try to get table data from motion file directly motion_data = motion_file_data[motion_key] if "table_pos" in motion_data and "table_quat" in motion_data: pos = torch.tensor(motion_data["table_pos"], device=self.device).float() quat = torch.tensor( motion_data["table_quat"], device=self.device ).float() scale = torch.tensor( motion_data.get("table_scale", [1.0, 1.0, 1.0]), device=self.device ).float() table_pos_list.append(pos) table_quat_list.append(quat) table_scale_list.append(scale) loaded_from_motion_file += 1 else: # No table data in motion file, use default table_pos_list.append( torch.tensor([0.0, 0.0, 0.8], device=self.device).float() ) table_quat_list.append( torch.tensor([1.0, 0.0, 0.0, 0.0], device=self.device).float() ) table_scale_list.append( torch.tensor([1.0, 1.0, 1.0], device=self.device).float() ) else: # Fallback: use default table position table_pos_list.append( torch.tensor([0.0, 0.0, 0.8], device=self.device).float() ) table_quat_list.append( torch.tensor([1.0, 0.0, 0.0, 0.0], device=self.device).float() ) table_scale_list.append( torch.tensor([1.0, 1.0, 1.0], device=self.device).float() ) self._replay_table_pos = torch.stack(table_pos_list, dim=0) # (num_envs, 3) self._replay_table_quat = torch.stack(table_quat_list, dim=0) # (num_envs, 4) self._replay_table_scale = torch.stack(table_scale_list, dim=0) # (num_envs, 3) # Apply fixed table_offset if configured table_offset = self.config.get("table_offset", None) if table_offset is not None: offset_tensor = torch.tensor( table_offset, device=self.device, dtype=self._replay_table_pos.dtype ) self._replay_table_pos = self._replay_table_pos + offset_tensor # Apply maximal X offset based on object starting position # This computes the maximum valid offset range such that object stays on table x_offset_maximal = self.config.get("replay_table_x_offset_maximal", False) x_offset_margin = self.config.get( "replay_table_x_offset_margin", 0.05 ) # 5cm safety margin if x_offset_maximal: # Get base table width from config (required for maximal offset mode) table_size_cfg = self.config.get("table_size") if table_size_cfg is None: raise ValueError( "replay_table_x_offset_maximal requires table_size to be set in config" ) base_table_width = table_size_cfg[0] # Get object starting position (step 0) for each env object_start_pos = self._motion_lib.get_object_root_pos( motion_ids, torch.zeros(num_replay_envs, device=self.device, dtype=torch.long), )[ :, 0 ] # Shape: (num_envs, 3) -> take first object if multiple object_start_x = object_start_pos[:, 0] # (num_envs,) # Get table position and scaled width per env table_x = self._replay_table_pos[:, 0] # (num_envs,) scale_x = self._replay_table_scale[:, 0] # (num_envs,) actual_width = base_table_width * scale_x # (num_envs,) half_width = actual_width / 2.0 # Compute valid offset range per env: # Object must stay on table after offset dx # dx_min = object_x - table_x - half_width + margin # dx_max = object_x - table_x + half_width - margin relative_obj_x = object_start_x - table_x dx_min = relative_obj_x - half_width + x_offset_margin dx_max = relative_obj_x + half_width - x_offset_margin # Sample X offset uniformly within valid range per env rand_vals = torch.rand(num_replay_envs, device=self.device) x_offsets = dx_min + rand_vals * (dx_max - dx_min) self._replay_table_pos[:, 0] += x_offsets logger.info("Replay table X offset (maximal mode):") logger.info( f" Base table width: {base_table_width:.3f}m, margin: {x_offset_margin:.3f}m" ) logger.info(f" dx_min range: [{dx_min.min():.3f}, {dx_min.max():.3f}]") logger.info(f" dx_max range: [{dx_max.min():.3f}, {dx_max.max():.3f}]") logger.info(f" Applied offsets: [{x_offsets.min():.3f}, {x_offsets.max():.3f}]") # Apply fixed XY offset range (legacy mode, used if maximal mode is disabled) # Format: [x_min, x_max, y_min, y_max] - per-env randomized xy_offset_range = self.config.get("replay_table_xy_offset_range", None) if xy_offset_range is not None and not x_offset_maximal: x_min, x_max, y_min, y_max = xy_offset_range # Sample X offset per environment x_offsets = torch.empty(num_replay_envs, device=self.device).uniform_( x_min, x_max ) self._replay_table_pos[:, 0] += x_offsets # Sample Y offset per environment (typically y_min = y_max for fixed offset) y_offsets = torch.empty(num_replay_envs, device=self.device).uniform_( y_min, y_max ) self._replay_table_pos[:, 1] += y_offsets logger.info( f"Replay table XY offset: X=[{x_min:.3f}, {x_max:.3f}], Y=[{y_min:.3f}, {y_max:.3f}]" ) if loaded_count > 0: logger.info( f"Loaded table metadata for {loaded_count}/{num_replay_envs} environments from {meta_dir}" ) elif loaded_from_motion_file > 0: logger.info( f"Loaded table metadata for {loaded_from_motion_file}/{num_replay_envs} environments from motion file" # noqa: E501 ) except Exception as e: # noqa: BLE001 logger.info(f"Warning: Could not load table metadata: {e}") # Setup keyboard controls for replay if not self.config.headless: self._setup_replay_keyboard() logger.info("Replay Controls:") logger.info(" G: Pause/Resume") logger.info(" B: Toggle Reverse Play") logger.info(" LEFT/RIGHT: Step backward/forward (when paused)") logger.info(" R: Restart from beginning") logger.info(" ESC: Exit replay") logger.info(" +/-: Increase/Decrease speed") logger.info("") # Return metadata results = [] for i, (mid, nsteps) in enumerate( zip(motion_ids.tolist(), num_steps_per_env.tolist(), strict=False) ): # Get motion key from motion library (curr_motion_keys is the active list) if hasattr(self._motion_lib, "curr_motion_keys") and mid < len( self._motion_lib.curr_motion_keys ): motion_key = self._motion_lib.curr_motion_keys[mid] elif hasattr(self._motion_lib, "motion_keys") and mid < len( self._motion_lib.motion_keys ): motion_key = self._motion_lib.motion_keys[mid] else: motion_key = f"motion_{mid}" results.append( { "env_id": i, "motion_id": mid, "motion_key": motion_key, "num_steps": nsteps, "duration": nsteps / 50.0, } ) return results def _setup_replay_keyboard(self): """Setup keyboard controls for replay mode""" # noqa: D415 try: if not hasattr(self, "keyboard_interface"): from isaaclab.devices.keyboard.se2_keyboard import Se2Keyboard, Se2KeyboardCfg cfg = Se2KeyboardCfg() self.keyboard_interface = Se2Keyboard(cfg) # Add replay-specific callbacks self.keyboard_interface.add_callback("G", self._toggle_replay_pause) self.keyboard_interface.add_callback("LEFT", self._replay_step_backward) self.keyboard_interface.add_callback("RIGHT", self._replay_step_forward) self.keyboard_interface.add_callback("R", self._restart_replay) self.keyboard_interface.add_callback("ESCAPE", self._exit_replay) self.keyboard_interface.add_callback("EQUAL", self._increase_replay_speed) self.keyboard_interface.add_callback("MINUS", self._decrease_replay_speed) self.keyboard_interface.add_callback("B", self._toggle_reverse_play) except Exception as e: # noqa: BLE001 logger.info(f"Could not setup replay keyboard controls: {e}") def _setup_contact_center_visualizer(self): """Setup visualization markers for per-hand contact centers during replay.""" if not VISUALIZATION_AVAILABLE: return has_left = hasattr(self._motion_lib, "_motion_object_contact_center_left") has_right = hasattr(self._motion_lib, "_motion_object_contact_center_right") if not has_left and not has_right: logger.info("No contact center available in motion library for visualization") return try: self._replay_contact_visualizers = {} if has_left: left_cfg = VisualizationMarkersCfg( prim_path="/Visuals/Replay/contact_center_left", markers={ "contact": sim_utils.SphereCfg( radius=0.03, visual_material=sim_utils.PreviewSurfaceCfg( diffuse_color=(0.0, 0.0, 1.0) # Blue for left hand ), ), }, ) self._replay_contact_visualizers["left_hand"] = VisualizationMarkers(left_cfg) self._replay_contact_visualizers["left_hand"].set_visibility(True) if has_right: right_cfg = VisualizationMarkersCfg( prim_path="/Visuals/Replay/contact_center_right", markers={ "contact": sim_utils.SphereCfg( radius=0.03, visual_material=sim_utils.PreviewSurfaceCfg( diffuse_color=(0.0, 1.0, 1.0) # Green for right hand ), ), }, ) self._replay_contact_visualizers["right_hand"] = VisualizationMarkers(right_cfg) self._replay_contact_visualizers["right_hand"].set_visibility(True) hands = list(self._replay_contact_visualizers.keys()) logger.info(f"Contact center visualizer initialized for: {', '.join(hands)}") # Keep _replay_contact_visualizer as a truthy check for the update loop self._replay_contact_visualizer = True except Exception as e: # noqa: BLE001 logger.info(f"Could not setup contact center visualizer: {e}") self._replay_contact_visualizer = None def _update_contact_center_visualization(self): """Update per-hand contact center visualization during replay.""" if not getattr(self, "_replay_contact_visualizers", None): return from gear_sonic.isaac_utils.rotations import quat_rotate # Get object pose (shared by both hands) object_root_pos = self._motion_lib.get_object_root_pos( self._replay_motion_ids, self._replay_time_steps )[ :, 0, : ] # (num_envs, 3) object_root_quat = self._motion_lib.get_object_root_quat( self._replay_motion_ids, self._replay_time_steps )[ :, 0, : ] # (num_envs, 4) # Environment origin offsets if hasattr(self, "_replay_custom_origins"): env_origins = self._replay_custom_origins[self._replay_env_ids] else: env_origins = self.env.scene.env_origins[self._replay_env_ids] # Hidden position for markers when no contact hidden_pos = torch.tensor([[0.0, 0.0, -1000.0]], device=self.device) for hand, visualizer in self._replay_contact_visualizers.items(): contact_center = self._motion_lib.get_object_contact_center( self._replay_motion_ids, self._replay_time_steps, hand=hand ) if contact_center is None: visualizer.visualize(translations=hidden_pos) continue # Move markers far away for envs with no contact (zero contact center) valid_mask = torch.norm(contact_center, dim=-1) > 1e-6 rotated = quat_rotate(object_root_quat, contact_center, w_last=False) world_center = rotated + object_root_pos + env_origins world_center[~valid_mask] = hidden_pos visualizer.visualize(translations=world_center) def _toggle_replay_pause(self): """Toggle pause state for replay""" # noqa: D415 if hasattr(self, "_replay_active") and self._replay_active: self._replay_paused = not self._replay_paused status = "PAUSED" if self._replay_paused else "PLAYING" max_frame = self._replay_time_steps.max().item() logger.info(f"Replay {status} (max frame {max_frame}/{self._replay_max_num_steps})") def _replay_step_backward(self): """Step backward one frame when paused""" # noqa: D415 if hasattr(self, "_replay_active") and self._replay_active and self._replay_paused: self._replay_time_steps = torch.clamp(self._replay_time_steps - 1, min=0) self._update_replay_frame() max_frame = self._replay_time_steps.max().item() logger.info(f"Frame {max_frame}/{self._replay_max_num_steps}") self.env.sim.render() def _replay_step_forward(self): """Step forward one frame when paused""" # noqa: D415 if hasattr(self, "_replay_active") and self._replay_active and self._replay_paused: self._replay_time_steps = torch.minimum( self._replay_time_steps + 1, self._replay_num_steps_per_env - 1 ) self._update_replay_frame() max_frame = self._replay_time_steps.max().item() logger.info(f"Frame {max_frame}/{self._replay_max_num_steps}") self.env.sim.render() def _restart_replay(self): """Restart replay from beginning""" # noqa: D415 if hasattr(self, "_replay_active") and self._replay_active: self._replay_time_steps.fill_(0) self._replay_paused = False self._update_replay_frame() logger.info("Replay restarted from beginning") self.env.sim.render() def _exit_replay(self): """Exit replay mode""" # noqa: D415 if hasattr(self, "_replay_active"): self._replay_active = False # Hide contact center visualizers for vis in getattr(self, "_replay_contact_visualizers", {}).values(): vis.set_visibility(False) logger.info("Exiting replay mode") self.env.sim.render() def _increase_replay_speed(self): """Increase replay speed""" # noqa: D415 if hasattr(self, "_replay_active") and self._replay_active: self._replay_speed = min(5.0, self._replay_speed * 1.25) logger.info(f"Replay speed: {self._replay_speed:.2f}x") def _decrease_replay_speed(self): """Decrease replay speed""" # noqa: D415 if hasattr(self, "_replay_active") and self._replay_active: self._replay_speed = max(0.1, self._replay_speed * 0.8) logger.info(f"Replay speed: {self._replay_speed:.2f}x") def _toggle_reverse_play(self): """Toggle reverse playback mode""" # noqa: D415 if hasattr(self, "_replay_active") and self._replay_active: self._replay_reverse = not self._replay_reverse direction = "REVERSE" if self._replay_reverse else "FORWARD" logger.info(f"Replay direction: {direction}") def _update_replay_frame(self): """Update all robot states to match the current replay frames""" # noqa: D415 if not hasattr(self, "_replay_active") or not self._replay_active: return # Get motion data at current time steps for all environments root_pos = self._motion_lib.get_root_pos_w(self._replay_motion_ids, self._replay_time_steps) root_quat = self._motion_lib.get_root_quat_w( self._replay_motion_ids, self._replay_time_steps ) root_lin_vel = self._motion_lib.get_root_lin_vel_w( self._replay_motion_ids, self._replay_time_steps ) root_ang_vel = self._motion_lib.get_root_ang_vel_w( self._replay_motion_ids, self._replay_time_steps ) motion_lib_joint_pos = self._motion_lib.get_dof_pos( self._replay_motion_ids, self._replay_time_steps ) motion_lib_joint_vel = self._motion_lib.get_dof_vel( self._replay_motion_ids, self._replay_time_steps ) # Handle DOF mismatch between motion library (e.g., 29 DOF) and robot (e.g., 43 DOF) robot_num_joints = self.motion_command.robot.num_joints motion_lib_num_dof = motion_lib_joint_pos.shape[-1] if robot_num_joints > motion_lib_num_dof and self._body_joint_indices is not None: # Robot has more DOFs than motion lib (e.g., 43 DOF robot with 29 DOF motion data) # Use body joint indices for proper mapping num_envs = motion_lib_joint_pos.shape[0] # Create full joint tensors with zeros for all DOFs joint_pos = torch.zeros( num_envs, robot_num_joints, device=self.device, dtype=motion_lib_joint_pos.dtype ) joint_vel = torch.zeros( num_envs, robot_num_joints, device=self.device, dtype=motion_lib_joint_vel.dtype ) # Map motion lib data to body joint indices (using G1_ISAACLab_ORDER mapping) joint_pos[:, self._body_joint_indices] = motion_lib_joint_pos joint_vel[:, self._body_joint_indices] = motion_lib_joint_vel # Use hand DOFs from motion lib if available, otherwise default to zero hand_dof_pos = self._motion_lib.get_hand_dof_pos( self._replay_motion_ids, self._replay_time_steps ) if hand_dof_pos is not None: # Hand DOFs are the last N joints (in Isaac order, not G1_HAND_JOINTS order) num_hand_dof = hand_dof_pos.shape[-1] joint_pos[:, -num_hand_dof:] = hand_dof_pos else: joint_pos = motion_lib_joint_pos joint_vel = motion_lib_joint_vel # Add environment origin offsets (use custom grid origins if set) if hasattr(self, "_replay_custom_origins"): root_pos = root_pos + self._replay_custom_origins[self._replay_env_ids] else: root_pos = root_pos + self.env.scene.env_origins[self._replay_env_ids] # Write state to simulation for all environments self.motion_command.robot.write_joint_state_to_sim( joint_pos, joint_vel, env_ids=self._replay_env_ids ) self.motion_command.robot.write_root_state_to_sim( torch.cat([root_pos, root_quat, root_lin_vel, root_ang_vel], dim=-1), env_ids=self._replay_env_ids, ) # Get object motion data from motion library if hasattr(self._motion_lib, "_motion_object_root_pos") and hasattr( self._motion_lib, "_motion_object_root_quat" ): object_root_pos = self._motion_lib.get_object_root_pos( self._replay_motion_ids, self._replay_time_steps ) object_root_quat = self._motion_lib.get_object_root_quat( self._replay_motion_ids, self._replay_time_steps ) # Add environment origin offsets to object position if hasattr(self, "_replay_custom_origins"): object_root_pos = object_root_pos + self._replay_custom_origins[ self._replay_env_ids ].unsqueeze(1) else: object_root_pos = object_root_pos + self.env.scene.env_origins[ self._replay_env_ids ].unsqueeze(1) # Write object state to simulation (handle multiple objects) # Shape: object_root_pos is (num_envs, max_num_objects, 3) # Shape: object_root_quat is (num_envs, max_num_objects, 4) for obj_idx in range(object_root_pos.shape[1]): obj_pos = object_root_pos[:, obj_idx, :] obj_quat = object_root_quat[:, obj_idx, :] object_root_pose = torch.cat([obj_pos, obj_quat], dim=-1) self.env.scene["object"].write_root_pose_to_sim( object_root_pose, env_ids=self._replay_env_ids ) if hasattr(self.env, "scene") and "table" in self.env.scene.rigid_objects: # Use pre-loaded per-motion table metadata (loaded during setup_replay_motion) # _replay_table_pos and _replay_table_quat have shape (num_envs, 3) and (num_envs, 4) if self._replay_table_pos is not None: # Use per-env table positions directly (already the right shape) table_pos = self._replay_table_pos.clone() table_quat = self._replay_table_quat.clone() else: # Fallback: derive from object position with hardcoded offset table_pos = self._motion_lib.get_object_root_pos( self._replay_motion_ids, torch.zeros_like(self._replay_time_steps) )[:, 0].clone() table_pos[:, 2] = 0.76 # Table height table_pos[:, 1] -= 0.15 table_quat = torch.tensor([[1.0, 0.0, 0.0, 0.0]], device=self.device).repeat( len(self._replay_env_ids), 1 ) # Apply table_offset if configured (for fallback path) table_offset = self.config.get("table_offset", None) if table_offset is not None: offset_tensor = torch.tensor( table_offset, device=self.device, dtype=table_pos.dtype ) table_pos = table_pos + offset_tensor # Add environment origin offsets (same as robot and object) if hasattr(self, "_replay_custom_origins"): table_pos = table_pos + self._replay_custom_origins[self._replay_env_ids] else: table_pos = table_pos + self.env.scene.env_origins[self._replay_env_ids] table_root_pose = torch.cat([table_pos, table_quat], dim=-1) self.env.scene["table"].write_root_pose_to_sim( table_root_pose, env_ids=self._replay_env_ids ) # Visualize contact points if enabled if ( hasattr(self, "_replay_contact_visualizer") and self._replay_contact_visualizer is not None and hasattr(self, "_replay_vis_enabled") and self._replay_vis_enabled ): self._update_contact_center_visualization() self.env.sim.forward() def step_replay(self): """Step the replay forward for all environments. Call this in a loop to animate the motions. Returns False when replay is complete. """ # noqa: D205 if not hasattr(self, "_replay_active") or not self._replay_active: return False # Update frames based on speed and direction if not self._replay_paused: if self._replay_reverse: # Playing in reverse self._replay_time_steps -= int(self._replay_speed) # Handle beginning of motions per environment at_start = self._replay_time_steps < 0 if at_start.any(): if self._replay_loop: # Loop to end when reaching start self._replay_time_steps[at_start] = ( self._replay_num_steps_per_env[at_start] - 1 ).to(torch.long) else: # If any motion reached start and not looping, end replay logger.info("Replay complete (reversed to start)!") self._replay_active = False return False else: # Playing forward (normal) self._replay_time_steps += int(self._replay_speed) # Handle end of motions per environment completed = self._replay_time_steps >= self._replay_num_steps_per_env if completed.any(): if self._replay_loop: # Loop completed motions self._replay_time_steps[completed] = 0 # Decrement loop counter if it's an integer (countdown mode) if not isinstance(self._replay_loop, bool): self._replay_loop -= 1 logger.info(f"Loop completed, {self._replay_loop} loops remaining") if self._replay_loop <= 0: logger.info("All loops complete!") self._replay_active = False return False else: # If any motion is complete and not looping, end replay logger.info("Replay complete!") self._replay_active = False return False # Update the robot states self._update_replay_frame() # Print progress every 50 frames (based on max time step) max_time_step = self._replay_time_steps.max().item() if max_time_step % 50 == 0 and max_time_step > 0: progress = (max_time_step / self._replay_max_num_steps) * 100 logger.info( f"Progress: {progress:.1f}% (max frame {max_time_step}/{self._replay_max_num_steps})" ) return True