scripts/play_dropbear_live.py

#!/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__")