#!/usr/bin/env python3 """Run Dropbear playback with an A100-to-desktop live frame bridge. Isaac Sim renders RGB frames on the CUDA device. This adapter replaces the stock RecordVideo wrapper with an ffplay-backed window, allowing a display-only GPU to present those frames without asking it to run Isaac Sim or PhysX. """ from __future__ import annotations import math import os import runpy import subprocess import sys import time from collections import deque from pathlib import Path from typing import Any import gymnasium as gym import cv2 import numpy as np import torch import warp as wp def _early_pop_float_pair_arg(name: str) -> tuple[float, float] | None: """Consume an environment-shaping pair before task registration.""" if name not in sys.argv: return None index = sys.argv.index(name) try: low = float(sys.argv[index + 1]) high = float(sys.argv[index + 2]) except (IndexError, ValueError) as exc: raise ValueError(f"{name} requires two numbers") from exc del sys.argv[index : index + 3] return low, high def _early_pop_float_arg(name: str) -> float | None: """Consume an environment-shaping scalar before task registration.""" if name not in sys.argv: return None index = sys.argv.index(name) try: value = float(sys.argv[index + 1]) except (IndexError, ValueError) as exc: raise ValueError(f"{name} requires a number") from exc del sys.argv[index : index + 2] return value ACTOR_BASE_LIN_VEL = "--actor-base-lin-vel" in sys.argv if ACTOR_BASE_LIN_VEL: sys.argv.remove("--actor-base-lin-vel") os.environ["DROPBEAR_ACTOR_BASE_LIN_VEL"] = "1" RESET_JOINT_POSITION_RANGE = _early_pop_float_pair_arg( "--viewer-reset-joint-position-range" ) RESET_JOINT_VELOCITY_RANGE = _early_pop_float_pair_arg( "--viewer-reset-joint-velocity-range" ) if RESET_JOINT_POSITION_RANGE is not None: reset_min, reset_max = RESET_JOINT_POSITION_RANGE if not 0.1 <= reset_min <= reset_max <= 2.0: raise ValueError( "--viewer-reset-joint-position-range must satisfy " "0.1 <= low <= high <= 2.0" ) os.environ["DROPBEAR_RESET_JOINT_POSITION_MIN"] = f"{reset_min:g}" os.environ["DROPBEAR_RESET_JOINT_POSITION_MAX"] = f"{reset_max:g}" if RESET_JOINT_VELOCITY_RANGE is not None: reset_vel_min, reset_vel_max = RESET_JOINT_VELOCITY_RANGE if not -5.0 <= reset_vel_min <= reset_vel_max <= 5.0: raise ValueError( "--viewer-reset-joint-velocity-range must satisfy " "-5.0 <= low <= high <= 5.0" ) os.environ["DROPBEAR_RESET_JOINT_VELOCITY_MIN"] = f"{reset_vel_min:g}" os.environ["DROPBEAR_RESET_JOINT_VELOCITY_MAX"] = f"{reset_vel_max:g}" RESET_POLICY_JOINTS_ONLY = "--viewer-reset-policy-joints-only" in sys.argv if RESET_POLICY_JOINTS_ONLY: sys.argv.remove("--viewer-reset-policy-joints-only") os.environ["DROPBEAR_RESET_POLICY_JOINTS_ONLY"] = "1" VIEWER_PLANE_ONLY = "--viewer-plane" in sys.argv if VIEWER_PLANE_ONLY: sys.argv.remove("--viewer-plane") os.environ["DROPBEAR_PLANE_ONLY"] = "1" VIEWER_OBSTACLE_TERRAIN = "--viewer-obstacle-terrain" in sys.argv if VIEWER_OBSTACLE_TERRAIN: sys.argv.remove("--viewer-obstacle-terrain") VIEWER_DIRECTIONAL_OBSTACLE_TERRAIN = ( "--viewer-directional-obstacle-terrain" in sys.argv ) if VIEWER_DIRECTIONAL_OBSTACLE_TERRAIN: sys.argv.remove("--viewer-directional-obstacle-terrain") if VIEWER_OBSTACLE_TERRAIN and VIEWER_DIRECTIONAL_OBSTACLE_TERRAIN: raise ValueError( "--viewer-obstacle-terrain conflicts with " "--viewer-directional-obstacle-terrain" ) if VIEWER_OBSTACLE_TERRAIN or VIEWER_DIRECTIONAL_OBSTACLE_TERRAIN: if VIEWER_PLANE_ONLY: raise ValueError( "Viewer obstacle terrain conflicts with --viewer-plane" ) if VIEWER_DIRECTIONAL_OBSTACLE_TERRAIN: os.environ["DROPBEAR_DIRECTIONAL_OBSTACLE_TERRAIN"] = "1" os.environ["DROPBEAR_DIRECTIONAL_OBSTACLE_COMMANDS"] = "1" else: os.environ["DROPBEAR_OBSTACLE_TERRAIN"] = "1" os.environ["DROPBEAR_TERRAIN_HEIGHT_SCAN"] = "1" VIEWER_DISABLE_PUSHES = "--viewer-disable-pushes" in sys.argv if VIEWER_DISABLE_PUSHES: sys.argv.remove("--viewer-disable-pushes") os.environ["DROPBEAR_DISABLE_PUSHES"] = "1" VIEWER_POSE_SEQUENCE = "--viewer-pose-sequence" in sys.argv if VIEWER_POSE_SEQUENCE: sys.argv.remove("--viewer-pose-sequence") os.environ["DROPBEAR_POSE_SEQUENCE"] = "1" VIEWER_COM_CONTROL = "--viewer-com-control" in sys.argv if VIEWER_COM_CONTROL: sys.argv.remove("--viewer-com-control") os.environ["DROPBEAR_COM_CONTROL"] = "1" VIEWER_GAIT_PERIOD = _early_pop_float_arg("--viewer-gait-period") if VIEWER_GAIT_PERIOD is not None: if not 0.20 <= VIEWER_GAIT_PERIOD <= 2.0: raise ValueError( "--viewer-gait-period must be between 0.20 and 2.0 seconds" ) os.environ["DROPBEAR_GAIT_PERIOD"] = f"{VIEWER_GAIT_PERIOD:g}" for cli_name, env_name, minimum, maximum in ( ("--viewer-com-stand-height", "DROPBEAR_COM_STAND_HEIGHT", 0.0, 3.0), ("--viewer-com-height-delta", "DROPBEAR_COM_HEIGHT_DELTA", 0.0, 1.0), ( "--viewer-com-height-error-scale", "DROPBEAR_COM_HEIGHT_ERROR_SCALE", 0.001, 1.0, ), ( "--viewer-com-vertical-velocity-error-scale", "DROPBEAR_COM_VERTICAL_VELOCITY_ERROR_SCALE", 0.001, 5.0, ), ): value = _early_pop_float_arg(cli_name) if value is not None: if not minimum <= value <= maximum: raise ValueError( f"{cli_name} must be between {minimum:g} and {maximum:g}" ) os.environ[env_name] = f"{value:g}" VIEWER_POSE_RESIDUAL_SCALE = _early_pop_float_arg("--viewer-pose-residual-scale") VIEWER_POSE_PG_PITCH_OFFSET = _early_pop_float_arg( "--viewer-pose-pg-pitch-offset" ) VIEWER_POSE_KNEE_OFFSET = _early_pop_float_arg("--viewer-pose-knee-offset") VIEWER_POSE_ANKLE67_OFFSET = _early_pop_float_arg( "--viewer-pose-ankle67-offset" ) for value, env_name, cli_name in ( ( VIEWER_POSE_PG_PITCH_OFFSET, "DROPBEAR_POSE_PG_PITCH_OFFSET", "--viewer-pose-pg-pitch-offset", ), ( VIEWER_POSE_KNEE_OFFSET, "DROPBEAR_POSE_KNEE_OFFSET", "--viewer-pose-knee-offset", ), ( VIEWER_POSE_ANKLE67_OFFSET, "DROPBEAR_POSE_ANKLE67_OFFSET", "--viewer-pose-ankle67-offset", ), ): if value is not None: if not -2.0 <= value <= 2.0: raise ValueError(f"{cli_name} must be between -2 and 2") os.environ[env_name] = f"{value:g}" VIEWER_POSE_BASELINE_DEPTH = _early_pop_float_arg("--viewer-pose-baseline-depth") if VIEWER_POSE_BASELINE_DEPTH is None: VIEWER_POSE_BASELINE_DEPTH = -0.33 if not -2.0 <= VIEWER_POSE_BASELINE_DEPTH <= 2.0: raise ValueError("--viewer-pose-baseline-depth must be between -2 and 2") os.environ["DROPBEAR_POSE_BASELINE_DEPTH"] = f"{VIEWER_POSE_BASELINE_DEPTH:g}" VIEWER_POSE_DEPTH_AMPLITUDE = _early_pop_float_arg( "--viewer-pose-depth-amplitude" ) if VIEWER_POSE_DEPTH_AMPLITUDE is None: VIEWER_POSE_DEPTH_AMPLITUDE = 1.0 VIEWER_POSE_STAND_HEIGHT = _early_pop_float_arg("--viewer-pose-stand-height") if VIEWER_POSE_STAND_HEIGHT is None: VIEWER_POSE_STAND_HEIGHT = 1.62 VIEWER_POSE_CROUCH_HEIGHT_DELTA = _early_pop_float_arg( "--viewer-pose-crouch-height-delta" ) if VIEWER_POSE_CROUCH_HEIGHT_DELTA is None: VIEWER_POSE_CROUCH_HEIGHT_DELTA = 0.18 VIEWER_POSE_ACTION_RESIDUAL = "--viewer-pose-action-residual" in sys.argv if VIEWER_POSE_ACTION_RESIDUAL: sys.argv.remove("--viewer-pose-action-residual") if not VIEWER_POSE_SEQUENCE: raise ValueError( "--viewer-pose-action-residual requires --viewer-pose-sequence" ) os.environ["DROPBEAR_POSE_ACTION_RESIDUAL"] = "1" if VIEWER_POSE_RESIDUAL_SCALE is not None: if not 0.0 <= VIEWER_POSE_RESIDUAL_SCALE <= 1.0: raise ValueError( "--viewer-pose-residual-scale must be between 0 and 1" ) os.environ["DROPBEAR_POSE_RESIDUAL_SCALE"] = ( f"{VIEWER_POSE_RESIDUAL_SCALE:g}" ) elif VIEWER_POSE_RESIDUAL_SCALE is not None: raise ValueError( "--viewer-pose-residual-scale requires --viewer-pose-action-residual" ) VIEWER_RECIPROCAL_SHOULDER_ACTIONS = ( "--viewer-reciprocal-shoulder-actions" in sys.argv ) if VIEWER_RECIPROCAL_SHOULDER_ACTIONS: sys.argv.remove("--viewer-reciprocal-shoulder-actions") os.environ["DROPBEAR_RECIPROCAL_SHOULDER_ACTIONS"] = "1" VIEWER_SHOULDER_COUNTERWEIGHT_SCALE = _early_pop_float_arg( "--viewer-shoulder-counterweight-scale" ) if VIEWER_SHOULDER_COUNTERWEIGHT_SCALE is not None: if VIEWER_RECIPROCAL_SHOULDER_ACTIONS: raise ValueError( "--viewer-reciprocal-shoulder-actions conflicts with " "--viewer-shoulder-counterweight-scale" ) if not 0.0 <= VIEWER_SHOULDER_COUNTERWEIGHT_SCALE <= 1.0: raise ValueError( "--viewer-shoulder-counterweight-scale must be between 0 and 1" ) os.environ["DROPBEAR_SHOULDER_COUNTERWEIGHT_SCALE"] = ( f"{VIEWER_SHOULDER_COUNTERWEIGHT_SCALE:g}" ) os.environ["DROPBEAR_RECIPROCAL_SHOULDER_ACTIONS"] = "1" import dropbear_walk # noqa: F401 import isaaclab.sim as sim_utils from isaaclab.sensors import CameraCfg from isaaclab.utils.math import quat_apply_inverse, yaw_quat from dropbear_walk.isaaclab_asset.dropbear import DISTAL_FOOT_BODIES from dropbear_walk.mdp.observations import ( com_height_reference, mass_weighted_com_state, pose_sequence_reference, ) from dropbear_walk.pose_guidance import ConstrainedPoseGuide from dropbear_walk.live_training_bridge import read_training_snapshot from rsl_rl.runners import OnPolicyRunner WORKSPACE_ROOT = Path(__file__).resolve().parents[1] PLAY_SCRIPT = WORKSPACE_ROOT / "IsaacLab" / "scripts" / "reinforcement_learning" / "rsl_rl" / "play.py" LIVE_WINDOW_SCRIPT = WORKSPACE_ROOT / "scripts" / "dropbear_live_window.py" ORIGINAL_GYM_MAKE = gym.make GUIDED_SEQUENCE = "--guided-sequence" in sys.argv if GUIDED_SEQUENCE: sys.argv.remove("--guided-sequence") FOLLOW_STATE: dict[str, Any] = { "checkpoint": None, "actor_checkpoint": None, "iteration": None, "command": 0.10, "reward": None, "episode_length": None, "tracking": None, "fall_rate": None, "loss": None, "loss_history": [], } def _pop_path_arg(name: str) -> Path | None: if name not in sys.argv: return None index = sys.argv.index(name) try: value = Path(sys.argv[index + 1]).expanduser().resolve() except IndexError as exc: raise ValueError(f"{name} requires a directory") from exc del sys.argv[index : index + 2] return value def _pop_float_arg(name: str) -> float | None: if name not in sys.argv: return None index = sys.argv.index(name) try: value = float(sys.argv[index + 1]) except (IndexError, ValueError) as exc: raise ValueError(f"{name} requires a number") from exc del sys.argv[index : index + 2] return value def _pop_int_arg(name: str) -> int | None: if name not in sys.argv: return None index = sys.argv.index(name) try: value = int(sys.argv[index + 1]) except (IndexError, ValueError) as exc: raise ValueError(f"{name} requires an integer") from exc del sys.argv[index : index + 2] return value FOLLOW_CHECKPOINT_DIR = _pop_path_arg("--follow-checkpoints") PIN_MANUAL_POLICY = "--pin-manual-policy" in sys.argv if PIN_MANUAL_POLICY: sys.argv.remove("--pin-manual-policy") try: _checkpoint_arg_index = sys.argv.index("--checkpoint") PINNED_ACTOR_CHECKPOINT = Path( sys.argv[_checkpoint_arg_index + 1] ).stem except (ValueError, IndexError): PINNED_ACTOR_CHECKPOINT = "startup actor" TRAINING_LIVE_STATE = _pop_path_arg("--training-live-state") VIEWER_COMMAND = _pop_float_arg("--viewer-command") if VIEWER_COMMAND is not None and not -2.0 <= VIEWER_COMMAND <= 2.0: raise ValueError("--viewer-command must be between -2 and 2 m/s") VIEWER_LATERAL_SPEED = _pop_float_arg("--viewer-lateral-speed") if ( VIEWER_LATERAL_SPEED is not None and not -2.0 <= VIEWER_LATERAL_SPEED <= 2.0 ): raise ValueError("--viewer-lateral-speed must be between -2 and 2 m/s") VIEWER_YAW_RATE = _pop_float_arg("--viewer-yaw-rate") if VIEWER_YAW_RATE is not None and not -3.0 <= VIEWER_YAW_RATE <= 3.0: raise ValueError("--viewer-yaw-rate must be between -3 and 3 rad/s") VIEWER_JOYSTICK_MAX_SPEED = _pop_float_arg("--viewer-joystick-max-speed") if VIEWER_JOYSTICK_MAX_SPEED is None: VIEWER_JOYSTICK_MAX_SPEED = 1.0 if not 0.05 <= VIEWER_JOYSTICK_MAX_SPEED <= 2.0: raise ValueError( "--viewer-joystick-max-speed must be between 0.05 and 2.0 m/s" ) VIEWER_JOYSTICK_MAX_YAW_RATE = _pop_float_arg( "--viewer-joystick-max-yaw-rate" ) if VIEWER_JOYSTICK_MAX_YAW_RATE is None: VIEWER_JOYSTICK_MAX_YAW_RATE = 1.0 if not 0.05 <= VIEWER_JOYSTICK_MAX_YAW_RATE <= 3.0: raise ValueError( "--viewer-joystick-max-yaw-rate must be between 0.05 and 3.0 rad/s" ) VIEWER_FPS = _pop_float_arg("--viewer-fps") if VIEWER_FPS is None: VIEWER_FPS = 50.0 if not 0.25 <= VIEWER_FPS <= 50.0: raise ValueError("--viewer-fps must be between 0.25 and 50") VIEWER_DISPLAY_FPS = _pop_float_arg("--viewer-display-fps") if VIEWER_DISPLAY_FPS is None: VIEWER_DISPLAY_FPS = 50.0 if not 10.0 <= VIEWER_DISPLAY_FPS <= 60.0: raise ValueError("--viewer-display-fps must be between 10 and 60") VIEWER_RENDER_WIDTH = _pop_int_arg("--viewer-render-width") if VIEWER_RENDER_WIDTH is None: VIEWER_RENDER_WIDTH = 960 VIEWER_RENDER_HEIGHT = _pop_int_arg("--viewer-render-height") if VIEWER_RENDER_HEIGHT is None: VIEWER_RENDER_HEIGHT = 540 if not 320 <= VIEWER_RENDER_WIDTH <= 1920: raise ValueError("--viewer-render-width must be between 320 and 1920") if not 180 <= VIEWER_RENDER_HEIGHT <= 1080: raise ValueError("--viewer-render-height must be between 180 and 1080") VIEWER_PLANAR_DEMO = "--viewer-planar-demo" in sys.argv if VIEWER_PLANAR_DEMO: sys.argv.remove("--viewer-planar-demo") VIEWER_COMMAND_HOLD_S = _pop_float_arg("--viewer-command-hold-s") if VIEWER_COMMAND_HOLD_S is None: VIEWER_COMMAND_HOLD_S = 5.0 if not 1.0 <= VIEWER_COMMAND_HOLD_S <= 60.0: raise ValueError("--viewer-command-hold-s must be between 1 and 60 seconds") if VIEWER_PLANAR_DEMO and FOLLOW_CHECKPOINT_DIR is None: raise ValueError("--viewer-planar-demo requires --follow-checkpoints") def _checkpoint_iteration(path: Path) -> int: try: return int(path.stem.removeprefix("model_")) except ValueError: return -1 def _latest_checkpoint(directory: Path) -> Path | None: checkpoints = list(directory.glob("model_*.pt")) if not checkpoints: return None return max(checkpoints, key=_checkpoint_iteration) def _display_gpu_name() -> str: """Return the GPU currently driving the desktop display.""" try: result = subprocess.run( [ "nvidia-smi", "--query-gpu=name,display_active", "--format=csv,noheader", ], check=True, capture_output=True, text=True, timeout=2.0, ) for line in result.stdout.splitlines(): name, active = (part.strip() for part in line.rsplit(",", 1)) if active.casefold() == "enabled": return name except (OSError, subprocess.SubprocessError, ValueError): pass return "desktop GPU" def _cpu_name() -> str: """Return the host CPU model without hard-coding workstation hardware.""" try: for line in Path("/proc/cpuinfo").read_text().splitlines(): if line.startswith("model name"): return line.split(":", 1)[1].strip() except (OSError, IndexError): pass return "CPU" def _latest_training_state( directory: Path, max_step: int | None = None, ) -> dict[str, Any]: """Read compact progress metrics from the latest TensorBoard event file.""" from tensorboard.backend.event_processing.event_accumulator import EventAccumulator state: dict[str, Any] = { "iteration": max_step if max_step is not None and max_step >= 0 else None, "command": 0.10, "reward": None, "episode_length": None, "tracking": None, "fall_rate": None, "loss": None, "loss_history": [], } event_files = sorted(directory.glob("events.out*"), key=lambda path: path.stat().st_mtime) if not event_files: return state accumulator = EventAccumulator(str(event_files[-1]), size_guidance={"scalars": 0}) accumulator.Reload() scalar_tags = set(accumulator.Tags().get("scalars", [])) metric_tags = { "command": "Curriculum/lin_vel_cmd_levels", "reward": "Train/mean_reward", "episode_length": "Train/mean_episode_length", "tracking": "Episode_Reward/track_lin_vel_xy", "fall_rate": "Episode_Termination/bad_orientation", "loss": "Loss/value", } for label, tag in metric_tags.items(): if tag not in scalar_tags: continue events = accumulator.Scalars(tag) if max_step is not None and max_step >= 0: events = [event for event in events if event.step <= max_step] if events: state[label] = float(events[-1].value) if "Loss/value" in scalar_tags: loss_events = accumulator.Scalars("Loss/value") if max_step is not None and max_step >= 0: loss_events = [ event for event in loss_events if event.step <= max_step ] state["loss_history"] = [ (int(event.step), float(event.value)) for event in loss_events[-120:] ] if loss_events: state["iteration"] = int(loss_events[-1].step) return state if FOLLOW_CHECKPOINT_DIR is not None: if not FOLLOW_CHECKPOINT_DIR.is_dir(): raise NotADirectoryError(FOLLOW_CHECKPOINT_DIR) original_get_inference_policy = OnPolicyRunner.get_inference_policy class AutoReloadPolicy: """Inference facade that adopts each newly saved actor in-place.""" def __init__(self, runner, policy): self.runner = runner self.policy = policy self.last_checkpoint: Path | None = None self.call_count = 0 def _reload_if_needed(self) -> None: self.call_count += 1 if self.call_count != 1 and self.call_count % 100 != 0: return checkpoint = _latest_checkpoint(FOLLOW_CHECKPOINT_DIR) if checkpoint is None: return if checkpoint == self.last_checkpoint: return if not PIN_MANUAL_POLICY: self.runner.load( str(checkpoint), load_cfg={ "actor": True, "critic": False, "optimizer": False, "iteration": False, "rnd": False, }, ) self.runner.alg.eval_mode() self.last_checkpoint = checkpoint FOLLOW_STATE["checkpoint"] = checkpoint.stem FOLLOW_STATE["actor_checkpoint"] = ( PINNED_ACTOR_CHECKPOINT if PIN_MANUAL_POLICY else checkpoint.stem ) FOLLOW_STATE.update( _latest_training_state( FOLLOW_CHECKPOINT_DIR, _checkpoint_iteration(checkpoint), ) ) if VIEWER_COMMAND is not None: FOLLOW_STATE["command"] = VIEWER_COMMAND print( f"[LIVE FOLLOW] " f"{'Observed' if PIN_MANUAL_POLICY else 'Adopted'} " f"{checkpoint.name}" f"{f' with pinned {PINNED_ACTOR_CHECKPOINT}' if PIN_MANUAL_POLICY else ''} at " f"vx {FOLLOW_STATE['command']:+.2f} m/s, " f"vy {(VIEWER_LATERAL_SPEED or 0.0):+.2f} m/s, " f"yaw {(VIEWER_YAW_RATE or 0.0):+.2f} rad/s" ) def __call__(self, observations): self._reload_if_needed() return self.policy(observations) def reset(self, dones=None): return self.policy.reset(dones) def get_auto_reload_policy(self, device=None): policy = original_get_inference_policy(self, device=device) return AutoReloadPolicy(self, policy) OnPolicyRunner.get_inference_policy = get_auto_reload_policy def make_with_live_camera(env_id: str, *args: Any, **kwargs: Any) -> gym.Env: """Inject a camera sensor that exists independently of the Kit viewport.""" env_cfg = kwargs.get("cfg") if env_cfg is not None: if FOLLOW_CHECKPOINT_DIR is not None: # A progress viewer should test the command band learned by the # current partial policy, not Play-v0's final 1.0 m/s limit. ranges = env_cfg.commands.base_velocity.ranges forward_speed = ( VIEWER_COMMAND if VIEWER_COMMAND is not None else 0.10 ) ranges.lin_vel_x = (forward_speed, forward_speed) lateral_speed = VIEWER_LATERAL_SPEED or 0.0 ranges.lin_vel_y = (lateral_speed, lateral_speed) yaw_rate = VIEWER_YAW_RATE or 0.0 ranges.ang_vel_z = (yaw_rate, yaw_rate) if ( VIEWER_OBSTACLE_TERRAIN and not VIEWER_DIRECTIONAL_OBSTACLE_TERRAIN and TRAINING_LIVE_STATE is None ): # Guarantee the single displayed environment gets an anchored # box rather than landing in the flat retention fraction. generator = env_cfg.scene.terrain.terrain_generator generator.sub_terrains["flat"].proportion = 0.0 generator.sub_terrains["forward_box"].proportion = 1.0 generator.num_rows = 2 generator.num_cols = 1 env_cfg.scene.terrain.max_init_terrain_level = 1 elif ( VIEWER_DIRECTIONAL_OBSTACLE_TERRAIN and TRAINING_LIVE_STATE is not None ): # The trainer needs 64 terrain columns to distribute 512 # environments. The one-environment renderer only needs one # representative column for each label. Keeping all nine # difficulty rows preserves the obstacle-height curriculum # while cutting render-scene terrain geometry by about 9x. generator = env_cfg.scene.terrain.terrain_generator for sub_terrain in generator.sub_terrains.values(): sub_terrain.proportion = 1.0 generator.num_cols = len(generator.sub_terrains) env_cfg.scene.terrain.max_init_terrain_level = ( generator.num_rows - 1 ) elif ( VIEWER_OBSTACLE_TERRAIN and TRAINING_LIVE_STATE is not None ): # Flat plus one forward-box column is sufficient to mirror # every source level for the legacy obstacle curriculum. generator = env_cfg.scene.terrain.terrain_generator for sub_terrain in generator.sub_terrains.values(): sub_terrain.proportion = 1.0 generator.num_cols = len(generator.sub_terrains) env_cfg.scene.terrain.max_init_terrain_level = ( generator.num_rows - 1 ) elif not ( VIEWER_OBSTACLE_TERRAIN or VIEWER_DIRECTIONAL_OBSTACLE_TERRAIN ): # The finite play terrain can be walked off during a persistent # viewer. Use an infinite plane for legacy flat-ground follow. env_cfg.scene.terrain.terrain_type = "plane" env_cfg.scene.terrain.terrain_generator = None env_cfg.scene.terrain.max_init_terrain_level = 0 env_cfg.curriculum.terrain_levels = None policy_step_s = float(env_cfg.sim.dt) * int(env_cfg.decimation) presentation_stride = max( 1, round(1.0 / (VIEWER_FPS * policy_step_s)), ) presentation_period_s = presentation_stride * policy_step_s # Rendering every physics/policy step can make the simulation run # slower than wall time. Align rendering, camera readback, and frame # presentation at a lower cadence while leaving 200 Hz physics and # 50 Hz policy control untouched. env_cfg.sim.render_interval = int(env_cfg.decimation) * presentation_stride env_cfg.scene.live_camera = CameraCfg( prim_path="{ENV_REGEX_NS}/LiveCamera", update_period=presentation_period_s, data_types=["rgb"], width=VIEWER_RENDER_WIDTH, height=VIEWER_RENDER_HEIGHT, spawn=sim_utils.PinholeCameraCfg( focal_length=24.0, focus_distance=400.0, horizontal_aperture=20.955, clipping_range=(0.1, 100.0), ), # Fixed three-quarter view aimed at the origin. offset=CameraCfg.OffsetCfg( pos=(2.5, 2.5, 2.5), rot=(0.33985114, 0.82047325, -0.42470819, -0.17591989), convention="ros", ), ) return ORIGINAL_GYM_MAKE(env_id, *args, **kwargs) class SimulatorMotionTrails: """World-space footprint and COM history rendered as simulator geometry.""" _AGE_BINS = 8 _MAX_FOOTPRINTS = 36 _MAX_COM_POINTS = 100 def __init__(self, env: gym.Env, distal_foot_body_ids: list[int]): # Import after AppLauncher has started Isaac Sim. Importing the marker # module earlier loads pxr bindings before Kit and corrupts its schema # extension initialization. from isaaclab.markers import ( VisualizationMarkers, VisualizationMarkersCfg, ) self._env = env.unwrapped self._distal_foot_body_ids = distal_foot_body_ids self._footprints: deque[ tuple[np.ndarray, np.ndarray, int] ] = deque(maxlen=self._MAX_FOOTPRINTS) self._com_points: deque[np.ndarray] = deque( maxlen=self._MAX_COM_POINTS ) self._previous_contacts = np.zeros(2, dtype=bool) self._last_episode_length: int | None = None self._sample_counter = 0 contact_sensor = self._env.scene.sensors["contact_forces"] contact_ids, contact_names = contact_sensor.find_sensors( DISTAL_FOOT_BODIES, preserve_order=True, ) if contact_names != DISTAL_FOOT_BODIES: raise ValueError( "Live trails could not resolve distal-foot contact sensors: " f"{contact_names}" ) self._contact_sensor = contact_sensor self._contact_sensor_ids = contact_ids footprint_markers: dict[str, Any] = {} self._footprint_prototype_indices: dict[tuple[int, int], int] = {} side_colors = ( (0.10, 0.55, 1.00), # left: blue (1.00, 0.30, 0.06), # right: orange ) for side, base_color in enumerate(side_colors): for age_bin in range(self._AGE_BINS): freshness = 1.0 - age_bin / (self._AGE_BINS - 1) brightness = 0.26 + 0.74 * freshness color = tuple(channel * brightness for channel in base_color) name = f"{'left' if side == 0 else 'right'}_age_{age_bin}" self._footprint_prototype_indices[(side, age_bin)] = len( footprint_markers ) footprint_markers[name] = sim_utils.CuboidCfg( size=(0.17, 0.072, 0.007), visual_material=sim_utils.PreviewSurfaceCfg( diffuse_color=color, emissive_color=tuple( channel * 0.12 for channel in color ), roughness=0.82, opacity=0.18 + 0.72 * freshness, ), ) self._footprint_markers = VisualizationMarkers( VisualizationMarkersCfg( prim_path="/World/DropbearMotionTrails/Footprints", markers=footprint_markers, ) ) self._footprint_markers.set_visibility(False) com_markers: dict[str, Any] = {} for age_bin in range(self._AGE_BINS): freshness = 1.0 - age_bin / (self._AGE_BINS - 1) color = ( 0.08 + 0.08 * freshness, 0.34 + 0.62 * freshness, 0.30 + 0.70 * freshness, ) com_markers[f"com_age_{age_bin}"] = sim_utils.CylinderCfg( radius=0.009, height=1.0, axis="X", visual_material=sim_utils.PreviewSurfaceCfg( diffuse_color=color, emissive_color=tuple( channel * 0.18 for channel in color ), roughness=0.65, opacity=0.16 + 0.76 * freshness, ), ) self._com_markers = VisualizationMarkers( VisualizationMarkersCfg( prim_path="/World/DropbearMotionTrails/CenterOfMassPath", markers=com_markers, ) ) self._com_markers.set_visibility(False) @staticmethod def _torch(value: Any) -> torch.Tensor: return wp.to_torch(value) if isinstance(value, wp.array) else value @staticmethod def _x_axis_to_vector_quaternion(vector: np.ndarray) -> np.ndarray: """Return an xyzw quaternion rotating +X onto ``vector``.""" direction = vector / max(float(np.linalg.norm(vector)), 1.0e-9) if direction[0] < -0.999999: return np.asarray((0.0, 0.0, 1.0, 0.0), dtype=np.float32) quaternion = np.asarray( (0.0, -direction[2], direction[1], 1.0 + direction[0]), dtype=np.float64, ) quaternion /= max(float(np.linalg.norm(quaternion)), 1.0e-9) return quaternion.astype(np.float32) def clear(self) -> None: self._footprints.clear() self._com_points.clear() self._previous_contacts[:] = False self._footprint_markers.set_visibility(False) self._com_markers.set_visibility(False) def _update_footprint_geometry(self) -> None: if not self._footprints: self._footprint_markers.set_visibility(False) return count = len(self._footprints) translations = np.stack( [footprint[0] for footprint in self._footprints] ).astype(np.float32) orientations = np.stack( [footprint[1] for footprint in self._footprints] ).astype(np.float32) marker_indices: list[int] = [] scales: list[tuple[float, float, float]] = [] for index, (_, _, side) in enumerate(self._footprints): normalized_age = (count - 1 - index) / max( self._MAX_FOOTPRINTS - 1, 1, ) age_bin = min( self._AGE_BINS - 1, round(normalized_age * (self._AGE_BINS - 1)), ) marker_indices.append( self._footprint_prototype_indices[(side, age_bin)] ) footprint_scale = 1.0 - 0.18 * normalized_age scales.append((footprint_scale, footprint_scale, 1.0)) self._footprint_markers.set_visibility(True) self._footprint_markers.visualize( translations=translations, orientations=orientations, scales=np.asarray(scales, dtype=np.float32), marker_indices=marker_indices, ) def _update_com_geometry(self) -> None: if len(self._com_points) < 2: self._com_markers.set_visibility(False) return points = np.stack(self._com_points).astype(np.float32) vectors = points[1:] - points[:-1] lengths = np.linalg.norm(vectors, axis=1) valid = lengths > 1.0e-5 if not bool(valid.any()): self._com_markers.set_visibility(False) return vectors = vectors[valid] lengths = lengths[valid] starts = points[:-1][valid] ends = points[1:][valid] translations = (starts + ends) * 0.5 orientations = np.stack( [ self._x_axis_to_vector_quaternion(vector) for vector in vectors ] ) count = len(vectors) marker_indices: list[int] = [] scales: list[tuple[float, float, float]] = [] for index, length in enumerate(lengths): normalized_age = (count - 1 - index) / max( self._MAX_COM_POINTS - 2, 1, ) age_bin = min( self._AGE_BINS - 1, round(normalized_age * (self._AGE_BINS - 1)), ) marker_indices.append(age_bin) radius_scale = 1.0 - 0.42 * normalized_age scales.append((float(length), radius_scale, radius_scale)) self._com_markers.set_visibility(True) self._com_markers.visualize( translations=translations.astype(np.float32), orientations=orientations.astype(np.float32), scales=np.asarray(scales, dtype=np.float32), marker_indices=marker_indices, ) def update(self) -> None: episode_length = int(self._env.episode_length_buf[0]) if ( self._last_episode_length is not None and episode_length < self._last_episode_length ): self.clear() self._last_episode_length = episode_length robot = self._env.scene["robot"] body_pos_w = self._torch(robot.data.body_pos_w) root_quat_w = self._torch(robot.data.root_quat_w) env_origin = self._env.scene.env_origins[0].to(body_pos_w.device) contact_time = self._contact_sensor.data.current_contact_time if contact_time is not None: contact_time = self._torch(contact_time) contacts = ( contact_time[0, self._contact_sensor_ids] > 0.0 ).detach().cpu().numpy() new_contacts = contacts & ~self._previous_contacts if bool(new_contacts.any()): yaw_wxyz = yaw_quat(root_quat_w[0:1])[0] yaw_xyzw = yaw_wxyz[[1, 2, 3, 0]].detach().cpu().numpy() trail_changed = False for side in np.flatnonzero(new_contacts): position = ( body_pos_w[0, self._distal_foot_body_ids[int(side)]] .detach() .cpu() .numpy() .astype(np.float32) ) position[2] = float(env_origin[2]) + 0.006 duplicate = any( previous_side == int(side) and float(np.linalg.norm(previous_pos[:2] - position[:2])) < 0.035 for previous_pos, _, previous_side in reversed( self._footprints ) ) if duplicate: continue self._footprints.append( (position, yaw_xyzw.astype(np.float32), int(side)) ) trail_changed = True if trail_changed: self._update_footprint_geometry() self._previous_contacts = contacts self._sample_counter += 1 if self._sample_counter % 3 == 0: com_pos_w, _ = mass_weighted_com_state(self._env, "robot") com_position = ( self._torch(com_pos_w)[0].detach().cpu().numpy().astype(np.float32) ) if ( not self._com_points or float(np.linalg.norm(com_position - self._com_points[-1])) >= 0.006 ): self._com_points.append(com_position) self._update_com_geometry() class FfplayLiveViewer(gym.Wrapper): """Drop-in replacement for RecordVideo that presents every frame live.""" def __init__(self, env: gym.Env, **_: Any): super().__init__(env) self._viewer: subprocess.Popen[bytes] | None = None self._viewer_control_buffer = b"" self._frame_size: tuple[int, int] | None = None self._frame_index = 0 self._simulation_step_count = 0 self._wall_start_time = time.perf_counter() self._measurement_wall_start = self._wall_start_time self._measurement_sim_step_start = 0 self._measurement_frame_start = 0 self._presentation_stride = max( 1, round(1.0 / (VIEWER_FPS * float(env.unwrapped.step_dt))), ) self._warmup_frames = 0 self._camera_positioned = False self._camera_orbit_angle = math.radians(45.0) initial_horizontal_radius = math.hypot(3.6, 3.6) self._camera_orbit_distance = math.hypot( initial_horizontal_radius, 2.0, ) self._camera_elevation_angle = math.atan2( 2.0, initial_horizontal_radius, ) self._camera_orbit_step = math.radians(15.0) simulation_device = torch.device(env.unwrapped.device) if simulation_device.type == "cuda": simulation_index = simulation_device.index if simulation_index is None: simulation_index = torch.cuda.current_device() self._simulation_gpu_name = torch.cuda.get_device_name( simulation_index ) self._render_gpu_name = self._simulation_gpu_name else: self._simulation_gpu_name = _cpu_name() if torch.cuda.is_available(): self._render_gpu_name = torch.cuda.get_device_name( torch.cuda.current_device() ) else: self._render_gpu_name = self._simulation_gpu_name self._display_gpu_name = _display_gpu_name() self._speed_ema: float | None = None self._lateral_speed_ema: float | None = None self._command_step = 0 self._active_vx = float(FOLLOW_STATE["command"]) self._active_vy = VIEWER_LATERAL_SPEED or 0.0 self._active_yaw_rate = VIEWER_YAW_RATE or 0.0 self._active_command_label = "fixed" self._joystick_active = False self._joystick_vx = 0.0 self._joystick_vy = 0.0 self._joystick_yaw_rate = 0.0 self._training_passthrough = False self._training_previous: dict[str, Any] | None = None self._training_current: dict[str, Any] | None = None self._training_received_at = 0.0 self._training_interpolation_s = 0.20 self._training_state_mtime_ns = -1 self._training_last_error: str | None = None self._manual_snapshot: dict[str, Any] | None = None self._pose_guide = ConstrainedPoseGuide(env.unwrapped) if GUIDED_SEQUENCE else None self._pose_joint_ids: list[int] = [] self._pose_height_body_ids: list[int] = [] self._pose_offsets: torch.Tensor | None = None robot = env.unwrapped.scene["robot"] self._distal_foot_body_ids, matched_foot_names = robot.find_bodies( DISTAL_FOOT_BODIES, preserve_order=True, ) if matched_foot_names != DISTAL_FOOT_BODIES: raise ValueError( "Live viewer could not resolve the ordered distal feet: " f"{matched_foot_names}" ) self._knee_joint_ids, matched_knee_names = robot.find_joints( [ "LL_knee_actuator_joint", "RL_knee_actuator_joint", ], preserve_order=True, ) if matched_knee_names != [ "LL_knee_actuator_joint", "RL_knee_actuator_joint", ]: raise ValueError( "Live viewer could not resolve the ordered knees: " f"{matched_knee_names}" ) self._motion_trails = SimulatorMotionTrails( env, self._distal_foot_body_ids, ) self._feet_min_distance = float( os.environ.get("DROPBEAR_FEET_MIN_DISTANCE", "0.14") ) self._feet_min_lateral_separation = float( os.environ.get("DROPBEAR_FEET_MIN_LATERAL_SEPARATION", "0.08") ) if ( VIEWER_OBSTACLE_TERRAIN or VIEWER_DIRECTIONAL_OBSTACLE_TERRAIN ) and TRAINING_LIVE_STATE is not None: self._move_manual_view_to_obstacle_tile() if VIEWER_POSE_SEQUENCE: self._pose_joint_ids, _ = robot.find_joints( [ "PG_left_leg_pitch", "PG_right_leg_pitch", "LL_knee_actuator_joint", "RL_knee_actuator_joint", "LL_Revolute67", "RL_Revolute67", ], preserve_order=True, ) self._pose_offsets = torch.tensor( [-0.10, -0.10, 0.25, 0.25, -0.15, -0.15], device=env.unwrapped.device, ) self._pose_height_body_ids, _ = robot.find_bodies( [ "torso_RMD_X10__1_Rotor_1", "torso_RMD_X10Rotot_1", ], preserve_order=True, ) def _start_viewer(self, frame: np.ndarray) -> None: height, width = frame.shape[:2] self._frame_size = (width, height) viewer_env = os.environ.copy() command = [ sys.executable, str(LIVE_WINDOW_SCRIPT), "--width", str(width), "--height", str(height), "--display-width", "1200", "--joystick-max-speed", f"{VIEWER_JOYSTICK_MAX_SPEED:g}", "--joystick-max-yaw-rate", f"{VIEWER_JOYSTICK_MAX_YAW_RATE:g}", "--source-fps", f"{1.0 / (self._presentation_stride * float(self.env.unwrapped.step_dt)):g}", "--display-fps", f"{VIEWER_DISPLAY_FPS:g}", "--title", ( f"Dropbear LIVE — {self._simulation_gpu_name} simulation / " f"{self._render_gpu_name} render / " f"{self._display_gpu_name} display" ), ] if TRAINING_LIVE_STATE is not None: command.append("--training-live-enabled") self._viewer = subprocess.Popen( command, stdin=subprocess.PIPE, stdout=subprocess.PIPE, env=viewer_env, bufsize=0, ) if self._viewer.stdout is not None: os.set_blocking(self._viewer.stdout.fileno(), False) print( f"[LIVE VIEW] Opened interactive desktop viewer for " f"{width}x{height} {self._render_gpu_name}-rendered frames " f"(PID {self._viewer.pid})." ) def _poll_camera_controls(self) -> None: if self._viewer is None or self._viewer.stdout is None: return try: chunk = os.read(self._viewer.stdout.fileno(), 4096) except BlockingIOError: chunk = b"" if chunk: self._viewer_control_buffer += chunk while b"\n" in self._viewer_control_buffer: line, self._viewer_control_buffer = ( self._viewer_control_buffer.split(b"\n", 1) ) command = line.decode("utf-8", errors="replace").strip() if command.startswith("view_mode "): requested_mode = command.removeprefix("view_mode ").strip() if requested_mode == "training": self._enter_training_passthrough() elif requested_mode == "manual": self._leave_training_passthrough() continue if command.startswith("robot_velocity "): if self._training_passthrough: continue parts = command.split() if len(parts) != 4: continue try: forward = max(-1.0, min(1.0, float(parts[1]))) left = max(-1.0, min(1.0, float(parts[2]))) yaw_left = max(-1.0, min(1.0, float(parts[3]))) except ValueError: continue self._joystick_active = True self._joystick_vx = forward * VIEWER_JOYSTICK_MAX_SPEED self._joystick_vy = left * VIEWER_JOYSTICK_MAX_SPEED self._joystick_yaw_rate = ( yaw_left * VIEWER_JOYSTICK_MAX_YAW_RATE ) continue if command == "camera_left": self._camera_orbit_angle += self._camera_orbit_step elif command == "camera_right": self._camera_orbit_angle -= self._camera_orbit_step elif command.startswith("camera_orbit_delta "): parts = command.split() if len(parts) != 3: continue try: azimuth_delta = math.radians(float(parts[1])) elevation_delta = math.radians(float(parts[2])) except ValueError: continue self._camera_orbit_angle += azimuth_delta self._camera_elevation_angle = min( math.radians(75.0), max( math.radians(5.0), self._camera_elevation_angle + elevation_delta, ), ) self._camera_orbit_angle = math.atan2( math.sin(self._camera_orbit_angle), math.cos(self._camera_orbit_angle), ) continue else: continue self._camera_orbit_angle = math.atan2( math.sin(self._camera_orbit_angle), math.cos(self._camera_orbit_angle), ) print( "[LIVE VIEW] Camera orbit " f"{math.degrees(self._camera_orbit_angle):+.0f} degrees." ) @staticmethod def _as_torch(value: Any) -> torch.Tensor: return wp.to_torch(value) if isinstance(value, wp.array) else value def _move_manual_view_to_obstacle_tile(self) -> None: """Keep manual mode on a box tile while retaining the full train mesh.""" terrain = self.env.unwrapped.scene.terrain if ( terrain is None or terrain.terrain_origins is None or not hasattr(terrain, "terrain_levels") or not hasattr(terrain, "terrain_types") ): return generator = terrain.cfg.terrain_generator if ( generator is None or "forward_box" not in generator.sub_terrains ): return old_origin = terrain.env_origins[0].clone() level = min(1, int(terrain.terrain_origins.shape[0]) - 1) # The compact directional atlas has one column per semantic label. # Selecting the middle numeric column can put the manual actor on a # lateral/turning obstacle instead of the forward box it learned. terrain_type = list(generator.sub_terrains).index("forward_box") terrain.terrain_levels[0] = level terrain.terrain_types[0] = terrain_type terrain.env_origins[0] = terrain.terrain_origins[level, terrain_type] robot = self.env.unwrapped.scene["robot"] root_pos = self._as_torch(robot.data.root_pos_w)[0:1].clone() root_quat = self._as_torch(robot.data.root_quat_w)[0:1].clone() root_pos += terrain.env_origins[0:1] - old_origin robot.write_root_pose_to_sim_index( root_pose=torch.cat((root_pos, root_quat), dim=-1), env_ids=torch.tensor([0], device=root_pos.device), ) def _capture_manual_snapshot(self) -> dict[str, Any]: unwrapped = self.env.unwrapped robot = unwrapped.scene["robot"] origin = unwrapped.scene.env_origins[0:1] terrain = unwrapped.scene.terrain terrain_level = None terrain_type = None if ( terrain is not None and hasattr(terrain, "terrain_levels") and hasattr(terrain, "terrain_types") ): terrain_level = int(terrain.terrain_levels[0]) terrain_type = int(terrain.terrain_types[0]) return { "root_position_local": ( self._as_torch(robot.data.root_pos_w)[0:1] - origin )[0].detach().cpu().tolist(), "root_quaternion_xyzw": self._as_torch( robot.data.root_quat_w )[0].detach().cpu().tolist(), "root_linear_velocity_world": self._as_torch( robot.data.root_lin_vel_w )[0].detach().cpu().tolist(), "root_angular_velocity_world": self._as_torch( robot.data.root_ang_vel_w )[0].detach().cpu().tolist(), "joint_names": list(robot.joint_names), "joint_position": self._as_torch( robot.data.joint_pos )[0].detach().cpu().tolist(), "joint_velocity": self._as_torch( robot.data.joint_vel )[0].detach().cpu().tolist(), "command": [ self._active_vx, self._active_vy, self._active_yaw_rate, ], "episode_step": int(unwrapped.episode_length_buf[0]), "terrain": { "level": terrain_level, "type": terrain_type, "label": "manual", }, } def _enter_training_passthrough(self) -> None: if TRAINING_LIVE_STATE is None: print("[TRAIN LIVE] No training state bridge was configured.") return if self._training_passthrough: return self._manual_snapshot = self._capture_manual_snapshot() self._training_passthrough = True self._joystick_active = False self._motion_trails.clear() print("[TRAIN LIVE] Viewer switched to actual training environment.") def _leave_training_passthrough(self) -> None: if not self._training_passthrough: return self._training_passthrough = False if self._manual_snapshot is not None: self._apply_snapshot(self._manual_snapshot) self._motion_trails.clear() print("[TRAIN LIVE] Viewer returned to manual policy control.") def _refresh_training_snapshot(self) -> None: if TRAINING_LIVE_STATE is None: return try: mtime_ns = TRAINING_LIVE_STATE.stat().st_mtime_ns except OSError: return if mtime_ns == self._training_state_mtime_ns: return snapshot = read_training_snapshot(TRAINING_LIVE_STATE) if snapshot is None: return robot = self.env.unwrapped.scene["robot"] if snapshot.get("joint_names") != list(robot.joint_names): error = "training/viewer joint ordering differs" if error != self._training_last_error: print(f"[TRAIN LIVE] Ignoring snapshot: {error}.") self._training_last_error = error return prior = self._training_current self._training_previous = prior self._training_current = snapshot self._training_state_mtime_ns = mtime_ns self._training_received_at = time.perf_counter() if prior is not None: source_period = ( float(snapshot["source_wall_time"]) - float(prior["source_wall_time"]) ) self._training_interpolation_s = min( 0.40, max(0.04, source_period), ) reset_or_tile_change = ( int(snapshot["episode_step"]) < int(prior["episode_step"]) or snapshot["terrain"] != prior["terrain"] ) if reset_or_tile_change: self._training_previous = None self._training_last_error = None @staticmethod def _normalized_quaternion_lerp( previous: list[float], current: list[float], alpha: float, ) -> list[float]: first = np.asarray(previous, dtype=np.float64) second = np.asarray(current, dtype=np.float64) if float(np.dot(first, second)) < 0.0: second = -second result = (1.0 - alpha) * first + alpha * second result /= max(float(np.linalg.norm(result)), 1.0e-12) return result.tolist() def _interpolated_training_snapshot(self) -> dict[str, Any] | None: current = self._training_current if current is None: return None previous = self._training_previous if previous is None: return current alpha = min( 1.0, max( 0.0, (time.perf_counter() - self._training_received_at) / self._training_interpolation_s, ), ) result = dict(current) for key in ( "root_position_local", "root_linear_velocity_world", "root_angular_velocity_world", "joint_position", "joint_velocity", ): first = np.asarray(previous[key], dtype=np.float64) second = np.asarray(current[key], dtype=np.float64) result[key] = ((1.0 - alpha) * first + alpha * second).tolist() result["root_quaternion_xyzw"] = self._normalized_quaternion_lerp( previous["root_quaternion_xyzw"], current["root_quaternion_xyzw"], alpha, ) return result def _apply_snapshot(self, snapshot: dict[str, Any]) -> None: unwrapped = self.env.unwrapped robot = unwrapped.scene["robot"] device = torch.device(unwrapped.device) dtype = self._as_torch(robot.data.joint_pos).dtype terrain_state = snapshot.get("terrain") or {} terrain = unwrapped.scene.terrain level = terrain_state.get("level") terrain_type = terrain_state.get("type") terrain_label = terrain_state.get("label") if ( terrain is not None and terrain.terrain_origins is not None and level is not None and terrain_type is not None ): generator = terrain.cfg.terrain_generator if ( generator is not None and terrain_label in generator.sub_terrains ): # Live training publishes the trainer's 64-column type index. # The viewer uses a compact one-column-per-label atlas, so map # by semantic label instead of clamping the source index. terrain_type = list(generator.sub_terrains).index( terrain_label ) level = max( 0, min(int(level), int(terrain.terrain_origins.shape[0]) - 1), ) terrain_type = max( 0, min( int(terrain_type), int(terrain.terrain_origins.shape[1]) - 1, ), ) terrain.terrain_levels[0] = level terrain.terrain_types[0] = terrain_type terrain.env_origins[0] = terrain.terrain_origins[level, terrain_type] origin = unwrapped.scene.env_origins[0:1] local_position = torch.tensor( snapshot["root_position_local"], device=device, dtype=dtype, ).unsqueeze(0) quaternion = torch.tensor( snapshot["root_quaternion_xyzw"], device=device, dtype=dtype, ).unsqueeze(0) root_velocity = torch.tensor( snapshot["root_linear_velocity_world"] + snapshot["root_angular_velocity_world"], device=device, dtype=dtype, ).unsqueeze(0) joint_position = torch.tensor( snapshot["joint_position"], device=device, dtype=dtype, ).unsqueeze(0) joint_velocity = torch.tensor( snapshot["joint_velocity"], device=device, dtype=dtype, ).unsqueeze(0) env_ids = torch.tensor([0], device=device) robot.write_root_pose_to_sim_index( root_pose=torch.cat((local_position + origin, quaternion), dim=-1), env_ids=env_ids, ) robot.write_root_velocity_to_sim_index( root_velocity=root_velocity, env_ids=env_ids, ) robot.write_joint_position_to_sim_index( position=joint_position, env_ids=env_ids, ) robot.write_joint_velocity_to_sim_index( velocity=joint_velocity, env_ids=env_ids, ) command = snapshot.get("command") or [0.0, 0.0, 0.0] command_term = unwrapped.command_manager.get_term("base_velocity") command_term.vel_command_b[0, :3] = torch.tensor( command[:3], device=device, dtype=command_term.vel_command_b.dtype, ) self._active_vx = float(command[0]) self._active_vy = float(command[1]) self._active_yaw_rate = float(command[2]) self._active_command_label = ( "train live" if self._training_passthrough else "manual" ) unwrapped.episode_length_buf[0] = int( snapshot.get("episode_step", 0) ) def _apply_training_passthrough(self) -> None: self._refresh_training_snapshot() snapshot = self._interpolated_training_snapshot() if snapshot is not None: self._apply_snapshot(snapshot) def _present(self) -> None: self._poll_camera_controls() live_camera = self.env.unwrapped.scene["live_camera"] robot = self.env.unwrapped.scene["robot"] # Terrain curricula can place env_0 far from world origin. Point the # camera using the robot's actual world pose and keep it as a chase cam. root_pos_w = robot.data.root_pos_w if isinstance(root_pos_w, wp.array): root_pos_w = wp.to_torch(root_pos_w) root = root_pos_w[0:1].clone() env_origin = self.env.unwrapped.scene.env_origins[0:1].to(root.device) root_offset = root - env_origin root_valid = bool( np.isfinite(root.detach().cpu().numpy()).all() and float(root_offset[:, :2].abs().max()) < 20.0 and -1.0 < float(root_offset[:, 2]) < 5.0 ) target = root if root_valid else env_origin.clone() target[:, 2] += 0.8 eyes = target.clone() horizontal_radius = self._camera_orbit_distance * math.cos( self._camera_elevation_angle ) eyes[:, 0] += horizontal_radius * math.cos( self._camera_orbit_angle ) eyes[:, 1] += horizontal_radius * math.sin( self._camera_orbit_angle ) eyes[:, 2] += self._camera_orbit_distance * math.sin( self._camera_elevation_angle ) live_camera.set_world_poses_from_view(eyes, target) self._motion_trails.update() if self._frame_index == 0 or self._frame_index % 250 == 0: root_values = tuple(round(float(value), 3) for value in root[0]) origin_values = tuple(round(float(value), 3) for value in env_origin[0]) print( f"[LIVE VIEW] chase root={root_values} origin={origin_values} " f"valid={root_valid}" ) # The first image was rendered before the initial chase pose was set. # Let one simulation step render from the new pose before presenting. if not self._camera_positioned: self._camera_positioned = True return frame = live_camera.data.output.get("rgb") if frame is None: return if hasattr(frame, "detach"): frame = frame[0].detach().cpu().numpy() elif frame.ndim == 4: frame = frame[0] frame = np.asarray(frame, dtype=np.uint8) if frame.ndim != 3 or frame.shape[2] < 3: raise RuntimeError(f"Unexpected live-view frame shape: {frame.shape}") frame = np.ascontiguousarray(frame[:, :, :3]) root_lin_vel_b = robot.data.root_lin_vel_b if isinstance(root_lin_vel_b, wp.array): root_lin_vel_b = wp.to_torch(root_lin_vel_b) root_ang_vel_b = robot.data.root_ang_vel_b if isinstance(root_ang_vel_b, wp.array): root_ang_vel_b = wp.to_torch(root_ang_vel_b) forward_speed = float(root_lin_vel_b[0, 0]) lateral_speed = float(root_lin_vel_b[0, 1]) self._speed_ema = ( forward_speed if self._speed_ema is None else 0.95 * self._speed_ema + 0.05 * forward_speed ) self._lateral_speed_ema = ( lateral_speed if self._lateral_speed_ema is None else 0.95 * self._lateral_speed_ema + 0.05 * lateral_speed ) planar_speed_ema = math.hypot( self._speed_ema, self._lateral_speed_ema, ) frame_max = int(frame.max()) frame_std = float(frame.std()) if frame_max <= 8 or frame_std <= 1.0: self._warmup_frames += 1 if self._warmup_frames == 1 or self._warmup_frames % 50 == 0: print( f"[LIVE VIEW] Dropped blank frame {self._warmup_frames} " f"(max={frame_max}, std={frame_std:.2f})." ) return # Upscale the cheaper camera render before drawing the HUD so text and # controls remain crisp in the desktop window. if frame.shape[:2] != (540, 960): frame = cv2.resize( frame, (960, 540), interpolation=cv2.INTER_LINEAR, ) root_quat_w = robot.data.root_quat_w if isinstance(root_quat_w, wp.array): root_quat_w = wp.to_torch(root_quat_w) body_pos_w = robot.data.body_pos_w if isinstance(body_pos_w, wp.array): body_pos_w = wp.to_torch(body_pos_w) foot_delta_w = ( body_pos_w[0, self._distal_foot_body_ids[0]] - body_pos_w[0, self._distal_foot_body_ids[1]] ).unsqueeze(0) foot_distance = float(torch.linalg.vector_norm(foot_delta_w[0])) foot_delta_yaw = quat_apply_inverse( yaw_quat(root_quat_w[0:1]), foot_delta_w, ) foot_lateral_separation = float(foot_delta_yaw[0, 1]) feet_clear = ( foot_distance >= self._feet_min_distance and foot_lateral_separation >= self._feet_min_lateral_separation ) joint_pos = robot.data.joint_pos if isinstance(joint_pos, wp.array): joint_pos = wp.to_torch(joint_pos) knee_left = float(joint_pos[0, self._knee_joint_ids[0]]) knee_right = float(joint_pos[0, self._knee_joint_ids[1]]) gravity_w = torch.zeros( 1, 3, dtype=root_quat_w.dtype, device=root_quat_w.device, ) gravity_w[:, 2] = -1.0 gravity_b = quat_apply_inverse(root_quat_w[0:1], gravity_w) torso_roll = float( torch.atan2(-gravity_b[0, 1], -gravity_b[0, 2]) ) gait_phase = ( float(self.env.unwrapped.episode_length_buf[0]) * float(self.env.unwrapped.step_dt) / float(os.environ.get("DROPBEAR_GAIT_PERIOD", "0.60")) ) % 1.0 left_phase = gait_phase right_phase = (gait_phase + 0.5) % 1.0 swing_leg = "L" if left_phase >= 0.55 else "R" checkpoint_name = FOLLOW_STATE["checkpoint"] if checkpoint_name is not None: display_command = self._active_vx left_lines: list[str] = [] if self._training_passthrough: snapshot = self._training_current left_lines.append("view TRAIN LIVE domain randomization ON") if snapshot is None: left_lines.append("state waiting for trainer") else: terrain_state = snapshot.get("terrain") or {} level = terrain_state.get("level") terrain_label = terrain_state.get("label", "unknown") level_text = "--" if level is None else str(level) domain = snapshot.get("domain_randomization") or {} state_age = max( 0.0, time.time() - float(snapshot.get("source_wall_time", time.time())), ) left_lines.extend( ( ( f"source env {int(snapshot['source_env'])} " f"sample age {state_age:.2f} s" ), ( f"terrain {terrain_label} " f"difficulty level {level_text}" ), ( "DR mass " f"{float(domain['total_mass_kg']):.2f} kg" if domain.get("total_mass_kg") is not None else "DR mass unavailable" ), ( "DR friction static/dynamic " f"{float(domain['static_friction_mean']):.2f}/" f"{float(domain['dynamic_friction_mean']):.2f}" if ( domain.get("static_friction_mean") is not None and domain.get("dynamic_friction_mean") is not None ) else "DR friction unavailable" ), ) ) else: left_lines.append("view MANUAL joystick control") if self._training_passthrough: pass elif self._pose_guide is not None: display_command = float( self._pose_guide.command_term.vel_command_b[0, 0] ) mode = self._pose_guide.last_mode or "initializing" left_lines.extend( ( "task pose guide", f"mode {mode}", f"cmd vx {display_command:+.3f} m/s", ) ) elif VIEWER_POSE_SEQUENCE: depth, direction = pose_sequence_reference(self.env.unwrapped) depth_value = float(depth[0]) direction_value = float(direction[0]) joint_pos = robot.data.joint_pos default_joint_pos = robot.data.default_joint_pos if isinstance(joint_pos, wp.array): joint_pos = wp.to_torch(joint_pos) if isinstance(default_joint_pos, wp.array): default_joint_pos = wp.to_torch(default_joint_pos) joint_delta = ( joint_pos[0, self._pose_joint_ids] - default_joint_pos[0, self._pose_joint_ids] ) achieved_depth = float( torch.sum(joint_delta * self._pose_offsets) / torch.sum(torch.square(self._pose_offsets)) ) torso_height = float( torch.mean(body_pos_w[0, self._pose_height_body_ids, 2]) - env_origin[0, 2] ) mode = ( "lowering" if direction_value > 0.5 else "rising" if direction_value < -0.5 else "crouch_hold" if depth_value > 0.5 else "stand" ) residual_scale = ( VIEWER_POSE_RESIDUAL_SCALE if VIEWER_POSE_RESIDUAL_SCALE is not None else 1.0 ) left_lines.extend( ( "task pose transfer", f"mode {mode}", f"residual gain {residual_scale:.2f}", f"cue depth {depth_value:.3f}", ( "depth target/actual " f"{VIEWER_POSE_BASELINE_DEPTH + VIEWER_POSE_DEPTH_AMPLITUDE * depth_value:.3f}/" f"{achieved_depth:.3f}" ), ( "height target/actual " f"{VIEWER_POSE_STAND_HEIGHT - VIEWER_POSE_CROUCH_HEIGHT_DELTA * depth_value:.3f}/" f"{torso_height:.3f} m" ), ) ) elif VIEWER_COM_CONTROL: com_pos_w, com_lin_vel_w = mass_weighted_com_state( self.env.unwrapped, "robot", ) com_lin_vel_yaw = quat_apply_inverse( yaw_quat(root_quat_w), com_lin_vel_w, ) target_height, target_vertical_velocity = com_height_reference( self.env.unwrapped, float(os.environ.get("DROPBEAR_COM_STAND_HEIGHT", "1.20")), float(os.environ.get("DROPBEAR_COM_HEIGHT_DELTA", "0.0")), 16.0, ) actual_height = float(com_pos_w[0, 2] - env_origin[0, 2]) left_lines.extend( ( "task COM trajectory", f"mode {self._active_command_label}", f"cmd vx {self._active_vx:+.3f} m/s", f"COM vx {float(com_lin_vel_yaw[0, 0]):+.3f} m/s", f"cmd vy {self._active_vy:+.3f} m/s", f"COM vy {float(com_lin_vel_yaw[0, 1]):+.3f} m/s", f"cmd yaw {self._active_yaw_rate:+.3f} rad/s", f"body yaw {float(root_ang_vel_b[0, 2]):+.3f} rad/s", f"COM z target {float(target_height[0]):.3f} m", f"COM z actual {actual_height:.3f} m", ( "COM vz target/actual " f"{float(target_vertical_velocity[0]):+.3f}/" f"{float(com_lin_vel_w[0, 2]):+.3f} m/s" ), ) ) else: left_lines.extend( ( "task live policy", f"cmd vx {display_command:+.3f} m/s", ) ) left_lines.extend( ( f"body vx {self._speed_ema:+.3f} m/s", f"body vy {self._lateral_speed_ema:+.3f} m/s", f"body speed {planar_speed_ema:.3f} m/s", ( f"knee L/R {knee_left:+.3f}/{knee_right:+.3f} rad " f"swing {swing_leg}" ), f"torso roll {torso_roll:+.3f} rad", f"feet distance {foot_distance:.3f} m", f"feet lateral L-R {foot_lateral_separation:+.3f} m", f"feet {'CLEAR' if feet_clear else 'CROSSING / TOO CLOSE'}", ) ) if all( FOLLOW_STATE[name] is not None for name in ("reward", "episode_length", "tracking", "fall_rate") ): left_lines.extend( ( f"reward {FOLLOW_STATE['reward']:.2f}", f"episode length {FOLLOW_STATE['episode_length']:.1f}", f"tracking {FOLLOW_STATE['tracking']:.3f}", f"fall rate {FOLLOW_STATE['fall_rate']:.4f}", ) ) font = cv2.FONT_HERSHEY_SIMPLEX font_scale = 0.30 line_height = 13 left_x = 14 left_y = 18 shadow = (0, 0, 0) normal_color = (235, 245, 255) for index, line in enumerate(left_lines): color = normal_color if line.startswith("task "): color = (255, 255, 255) elif line.startswith("feet "): color = (120, 255, 150) if feet_clear else (255, 100, 90) y = left_y + index * line_height cv2.putText( frame, line, (left_x + 1, y + 1), font, font_scale, shadow, 2, cv2.LINE_AA, ) cv2.putText( frame, line, (left_x, y), font, font_scale, color, 1, cv2.LINE_AA, ) iteration = FOLLOW_STATE.get("iteration") if iteration is None: iteration = _checkpoint_iteration(Path(checkpoint_name)) viewer_elapsed_s = max( time.perf_counter() - self._measurement_wall_start, 1.0e-6, ) sim_realtime_factor = ( ( self._simulation_step_count - self._measurement_sim_step_start ) * float(self.env.unwrapped.step_dt) / viewer_elapsed_s ) achieved_render_fps = ( self._frame_index - self._measurement_frame_start ) / viewer_elapsed_s right_lines = [ f"epoch {int(iteration)}" if int(iteration) >= 0 else "epoch --", checkpoint_name, ( f"manual actor {FOLLOW_STATE['actor_checkpoint']}" if FOLLOW_STATE.get("actor_checkpoint") != checkpoint_name else f"actor {checkpoint_name}" ), f"{self._simulation_gpu_name} simulation", f"{self._render_gpu_name} render", f"{self._display_gpu_name} display", ( f"sim realtime x{sim_realtime_factor:.2f}" f" rendered {achieved_render_fps:.2f}" f" / display {VIEWER_DISPLAY_FPS:.0f} fps" ), ] if FOLLOW_STATE.get("loss") is not None: right_lines.append(f"value loss {FOLLOW_STATE['loss']:.5f}") right_margin = 14 for index, line in enumerate(right_lines): (text_width, _), _ = cv2.getTextSize( line, font, font_scale, 1, ) x = frame.shape[1] - right_margin - text_width y = left_y + index * line_height cv2.putText( frame, line, (x + 1, y + 1), font, font_scale, shadow, 2, cv2.LINE_AA, ) cv2.putText( frame, line, (x, y), font, font_scale, (245, 245, 255), 1, cv2.LINE_AA, ) loss_history = FOLLOW_STATE.get("loss_history") or [] if len(loss_history) >= 2: graph_width = min(285, frame.shape[1] // 3) graph_height = min(112, frame.shape[0] // 4) graph_x = frame.shape[1] - graph_width - 14 graph_y = frame.shape[0] - graph_height - 14 panel = frame.copy() cv2.rectangle( panel, (graph_x, graph_y), (graph_x + graph_width, graph_y + graph_height), (5, 8, 14), -1, ) cv2.addWeighted(panel, 0.72, frame, 0.28, 0.0, frame) cv2.rectangle( frame, (graph_x, graph_y), (graph_x + graph_width, graph_y + graph_height), (150, 165, 190), 1, ) plot_left = graph_x + 9 plot_right = graph_x + graph_width - 8 plot_top = graph_y + 24 plot_bottom = graph_y + graph_height - 10 values = np.asarray( [max(abs(float(value)), 1.0e-8) for _, value in loss_history], dtype=np.float64, ) log_values = np.log10(values) smoothed_log_values = np.empty_like(log_values) smoothed_log_values[0] = log_values[0] for index in range(1, len(log_values)): smoothed_log_values[index] = ( 0.25 * log_values[index] + 0.75 * smoothed_log_values[index - 1] ) low = float(log_values.min()) high = float(log_values.max()) if high - low < 1.0e-6: high = low + 1.0 x_values = np.linspace( plot_left, plot_right, len(log_values), ) y_values = plot_bottom - ( (log_values - low) / (high - low) * (plot_bottom - plot_top) ) points = np.column_stack((x_values, y_values)).astype(np.int32) smoothed_y_values = plot_bottom - ( (smoothed_log_values - low) / (high - low) * (plot_bottom - plot_top) ) smoothed_points = np.column_stack( (x_values, smoothed_y_values) ).astype(np.int32) cv2.line( frame, (plot_left, plot_bottom), (plot_right, plot_bottom), (105, 120, 145), 1, ) cv2.polylines( frame, [points], False, (90, 150, 180), 1, cv2.LINE_AA, ) cv2.polylines( frame, [smoothed_points], False, (110, 255, 205), 2, cv2.LINE_AA, ) first_step = int(loss_history[0][0]) last_step = int(loss_history[-1][0]) graph_title = ( f"value loss log raw {values[-1]:.4f} " f"EMA {10.0 ** smoothed_log_values[-1]:.4f}" ) cv2.putText( frame, graph_title, (graph_x + 8, graph_y + 15), font, 0.28, (235, 245, 255), 1, cv2.LINE_AA, ) cv2.putText( frame, f"epoch {first_step}-{last_step}", (graph_x + 8, graph_y + graph_height - 2), font, 0.25, (180, 195, 215), 1, cv2.LINE_AA, ) self._frame_index += 1 # Replicator intentionally returns black frames while its render product # warms up. Do not create a misleading empty window from those frames. if self._viewer is None: print( f"[LIVE VIEW] First scene frame ready after {self._warmup_frames} warm-up frames " f"(max={frame_max}, std={frame_std:.2f})." ) if self._viewer is None: self._start_viewer(frame) if self._viewer is not None: self._measurement_wall_start = time.perf_counter() self._measurement_sim_step_start = self._simulation_step_count self._measurement_frame_start = self._frame_index if self._viewer is None or self._viewer.stdin is None: return if self._viewer.poll() is not None: raise KeyboardInterrupt try: self._viewer.stdin.write(frame.tobytes()) self._viewer.stdin.flush() except BrokenPipeError as exc: raise KeyboardInterrupt from exc def step(self, action): self._poll_camera_controls() if FOLLOW_CHECKPOINT_DIR is not None and not self._training_passthrough: if self._joystick_active: self._active_command_label = "joystick" self._active_vx = self._joystick_vx self._active_vy = self._joystick_vy self._active_yaw_rate = self._joystick_yaw_rate elif VIEWER_PLANAR_DEMO: forward_speed = abs( float( VIEWER_COMMAND if VIEWER_COMMAND is not None else FOLLOW_STATE["command"] ) ) lateral_speed = abs(VIEWER_LATERAL_SPEED or 0.10) commands = ( ("forward", forward_speed, 0.0, 0.0), ("left", 0.0, lateral_speed, 0.0), ("right", 0.0, -lateral_speed, 0.0), ( "backward", -min(lateral_speed, forward_speed), 0.0, 0.0, ), ) hold_steps = max( 1, round( VIEWER_COMMAND_HOLD_S / float(self.env.unwrapped.step_dt) ), ) command_index = (self._command_step // hold_steps) % len(commands) ( self._active_command_label, self._active_vx, self._active_vy, self._active_yaw_rate, ) = commands[command_index] else: self._active_command_label = "fixed" self._active_vx = float(FOLLOW_STATE["command"]) self._active_vy = VIEWER_LATERAL_SPEED or 0.0 self._active_yaw_rate = VIEWER_YAW_RATE or 0.0 command_term = self.env.unwrapped.command_manager.get_term("base_velocity") command_term.vel_command_b[:, 0] = self._active_vx command_term.vel_command_b[:, 1] = self._active_vy command_term.vel_command_b[:, 2] = self._active_yaw_rate self._command_step += 1 if self._pose_guide is not None: action = self._pose_guide.apply(action) result = self.env.step(action) if self._training_passthrough: self._apply_training_passthrough() self._simulation_step_count += 1 if self._simulation_step_count % self._presentation_stride == 0: self._present() else: # Controls remain responsive at the 50 Hz policy rate even when # RGB readback/presentation is intentionally decimated. self._poll_camera_controls() # The stock player's --real-time limiter budgets only the environment # call and can leave persistent time dilation from work immediately # outside that call. Pace against the visible viewer epoch instead: # sleep only while ahead, and never add delay when rendering is behind. if self._viewer is not None: simulated_elapsed_s = ( self._simulation_step_count - self._measurement_sim_step_start ) * float(self.env.unwrapped.step_dt) wall_elapsed_s = time.perf_counter() - self._measurement_wall_start sleep_s = simulated_elapsed_s - wall_elapsed_s if sleep_s > 0.0: time.sleep(sleep_s) return result def reset(self, **kwargs): result = self.env.reset(**kwargs) if self._training_passthrough: self._apply_training_passthrough() self._present() return result def close(self) -> None: if self._viewer is not None: if self._viewer.stdin is not None: self._viewer.stdin.close() if self._viewer.poll() is None: self._viewer.terminate() try: self._viewer.wait(timeout=2.0) except subprocess.TimeoutExpired: self._viewer.kill() self._viewer = None self.env.close() if not PLAY_SCRIPT.is_file(): raise FileNotFoundError(f"Isaac Lab player not found: {PLAY_SCRIPT}") # The stock player already has the correct checkpoint/config/policy logic. # Its video path enables offscreen RGB rendering; replacing RecordVideo turns # that same stream into a live desktop window and avoids writing a movie. gym.wrappers.RecordVideo = FfplayLiveViewer gym.make = make_with_live_camera if "--video" not in sys.argv: sys.argv.extend(["--video", "--video_length", "2147483647"]) sys.path.insert(0, str(PLAY_SCRIPT.parent)) runpy.run_path(str(PLAY_SCRIPT), run_name="__main__")