| from __future__ import annotations |
|
|
| import io |
| import json |
| import math |
| from dataclasses import dataclass |
| from pathlib import Path |
| from typing import Callable |
|
|
| import numpy as np |
|
|
| from .spec import ( |
| ACTION_DIM, |
| ALL_BEHAVIORS, |
| BEHAVIORS, |
| CONTROL_DECIMATION, |
| CONTROL_HZ, |
| DEFAULT_POSE, |
| RECOVERY_STABLE_S, |
| STATE_DIM, |
| TERMINAL_STABILIZATION_S, |
| BehaviorSpec, |
| ) |
|
|
|
|
| BALL_OFFSET_X = 0.09 |
| BALL_OFFSET_ABS_Y = 0.042 |
| BALL_RADIUS = 0.035 |
| |
| |
| |
| HEAD_CAMERA_QUAT = np.array( |
| [0.7010574, -0.0922960, 0.0922960, -0.7010574], dtype=np.float64 |
| ) |
| RECOVERY_INSTRUCTION = ( |
| "Recover from the fall and stand upright before doing anything else." |
| ) |
| BALANCE_INSTRUCTION = "Stand upright, hold balance, and remain still." |
| AUTONOMOUS_TASK_INSTRUCTIONS = { |
| "sitstand": "Perform one sit-down and rise cycle.", |
| "ground_pick": "Touch the ground once with your beak.", |
| "roller_crouch": "Crouch once while gliding on the rollers.", |
| "roulade": "Perform one forward roll.", |
| } |
|
|
| ControlProgram = Callable[[float], tuple[str, np.ndarray, str, str]] |
|
|
|
|
| @dataclass |
| class EpisodeResult: |
| images: np.ndarray |
| states: np.ndarray |
| observation_control_indices: np.ndarray |
| actions: np.ndarray |
| instruction: str |
| success: bool |
| metrics: dict[str, float | bool | str] |
| command: np.ndarray |
| subtask: str |
| observation_instructions: tuple[str, ...] |
| observation_phases: tuple[str, ...] |
| video_frames: np.ndarray | None = None |
|
|
|
|
| class SessionCache: |
| def __init__(self, policies_dir: Path): |
| self.policies_dir = policies_dir |
| self._sessions: dict[str, object] = {} |
|
|
| def get(self, filename: str): |
| import onnxruntime as ort |
|
|
| if filename not in self._sessions: |
| path = self.policies_dir / filename |
| if not path.exists(): |
| raise FileNotFoundError(path) |
| options = ort.SessionOptions() |
| options.intra_op_num_threads = 1 |
| options.inter_op_num_threads = 1 |
| session = ort.InferenceSession( |
| str(path), sess_options=options, providers=["CPUExecutionProvider"] |
| ) |
| in_shape = list(session.get_inputs()[0].shape) |
| out_shape = list(session.get_outputs()[0].shape) |
| if in_shape[-1] != STATE_DIM or out_shape[-1] != ACTION_DIM: |
| raise ValueError( |
| f"{filename}: expected {STATE_DIM}->{ACTION_DIM}, got {in_shape}->{out_shape}" |
| ) |
| self._sessions[filename] = session |
| return self._sessions[filename] |
|
|
|
|
| def _quat_rotate_inverse(quat: np.ndarray, vec: np.ndarray) -> np.ndarray: |
| w = quat[0] |
| xyz = quat[1:4] |
| t = np.cross(xyz, vec) * 2 |
| return vec - w * t + np.cross(xyz, t) |
|
|
|
|
| def _instruction(spec: BehaviorSpec, command: np.ndarray, ordinal: int, subtask: str) -> str: |
| variant = ordinal % 6 |
| if spec.kind == "recovery": |
| return ( |
| "Recover from the fall, stand upright, and hold balance.", |
| "Get back onto both feet and stabilize in a standing pose.", |
| "Stand up from the ground, then remain upright and still.", |
| "First recover from the fall, then balance on both feet.", |
| )[ordinal % 4] |
| if spec.name == "stand": |
| z, roll, pitch = command[9], command[10], command[11] |
| templates = ( |
| "Stand still with body height offset {z:+.3f} m, roll {r:+.2f} rad and pitch {p:+.2f} rad.", |
| "Hold a stable standing pose: height {z:+.3f} m, roll {r:+.2f}, pitch {p:+.2f}.", |
| "Keep balance while standing at body offsets z={z:+.3f}, roll={r:+.2f}, pitch={p:+.2f}.", |
| "Remain upright and stationary with z {z:+.3f} m, roll {r:+.2f}, pitch {p:+.2f}.", |
| "Balance in place at the requested body pose: z={z:+.3f}, r={r:+.2f}, p={p:+.2f}.", |
| "Stay standing and steady; body pose z {z:+.3f}, roll {r:+.2f}, pitch {p:+.2f}.", |
| ) |
| return templates[variant].format(z=z, r=roll, p=pitch) |
| if spec.name == "walking": |
| vx, vy, wz = command[:3] |
| templates = ( |
| "Walk with forward velocity {x:+.2f} m/s, lateral velocity {y:+.2f} m/s and yaw rate {w:+.2f} rad/s.", |
| "Move at vx={x:+.2f}, vy={y:+.2f} m/s while turning at {w:+.2f} rad/s.", |
| "Follow the walking command: forward {x:+.2f}, sideways {y:+.2f}, rotation {w:+.2f}.", |
| "Walk steadily with velocity command ({x:+.2f}, {y:+.2f}, {w:+.2f}).", |
| "Travel using vx {x:+.2f} m/s, vy {y:+.2f} m/s and wz {w:+.2f} rad/s.", |
| "Execute the commanded gait: x {x:+.2f}, y {y:+.2f}, yaw {w:+.2f}.", |
| ) |
| return templates[variant].format(x=vx, y=vy, w=wz) |
| if spec.name == "sitstand": |
| if subtask == "stand_up": |
| return ( |
| "Stand up gently from sitting and finish balanced.", |
| "Rise from the seated pose and stabilize upright.", |
| "Move from sitting to a balanced standing pose.", |
| )[ordinal % 3] |
| return ( |
| "Sit down gently, then stand back up and balance.", |
| "Lower into a seated pose, rise again, and finish upright.", |
| "Perform one sit-to-stand cycle and stabilize on both feet.", |
| )[ordinal % 3] |
| if spec.name == "ground_pick": |
| return ( |
| "Lower your beak to touch the ground, then stand back up.", |
| "Perform one ground-touch cycle and return upright.", |
| "Crouch to the floor with your head and rise again.", |
| )[ordinal % 3] |
| if spec.name in ("kick_left", "kick_right"): |
| side = spec.name.split("_")[1] |
| if ordinal % 2 == 0: |
| return "Kick the visible ball." |
| return ( |
| f"Kick the ball with the {side} foot.", |
| f"Use your {side} leg to kick the ball forward.", |
| f"Perform a {side}-foot ball kick.", |
| )[ordinal % 3] |
| if spec.name == "roller": |
| vx = command[0] |
| return ( |
| f"Skate straight with push command {vx:+.2f}.", |
| f"Roll forward on the skates at command {vx:+.2f}.", |
| f"Use the rollers to move straight, drive {vx:+.2f}.", |
| )[ordinal % 3] |
| if spec.name == "roller_crouch": |
| return ( |
| "Crouch while gliding on the rollers, then stand again.", |
| "Perform one roller crouch cycle and recover upright.", |
| "Lower into a skating crouch and rise while rolling.", |
| )[ordinal % 3] |
| return ( |
| "Perform a forward roll and recover to standing.", |
| "Roll forward over your head, then stand up.", |
| "Do one forward roulade and finish upright.", |
| )[ordinal % 3] |
|
|
|
|
| class MicroDuckEpisode: |
| def __init__( |
| self, |
| spec: BehaviorSpec, |
| microduck_rl: Path, |
| sessions: SessionCache, |
| seed: int, |
| ordinal: int, |
| image_size: int = 256, |
| randomize: bool = True, |
| start_fallen: bool = False, |
| control_program: ControlProgram | None = None, |
| instruction_override: str | None = None, |
| ball_forward_offset: float | None = None, |
| ): |
| import mujoco |
|
|
| self.spec = spec |
| self.sessions = sessions |
| self.rng = np.random.default_rng(seed) |
| self.ordinal = ordinal |
| self.image_size = image_size |
| self.control_program = control_program |
| self.ball_forward_offset = ball_forward_offset |
| self.started_fallen = bool(start_fallen or spec.kind == "recovery") |
| self.recovery_active = self.started_fallen |
| self.recovery_completed = False |
| self.recovery_enter_step = 0 if self.started_fallen else None |
| self.recovery_exit_step: int | None = None |
| self.recovery_upright_ticks = 0 |
| self.task_start_step: int | None = None if self.started_fallen else 0 |
| self.task_control_steps = 0 |
| self.fall_direction = "" |
| xml = ( |
| microduck_rl |
| / "src/mjlab_microduck/robot/microduck" |
| / spec.scene_file |
| ) |
| if not xml.exists(): |
| raise FileNotFoundError(xml) |
| self.model = mujoco.MjModel.from_xml_path(str(xml)) |
| self.model.opt.timestep = 1.0 / 200.0 |
| self.data = mujoco.MjData(self.model) |
| self.renderer = mujoco.Renderer(self.model, height=image_size, width=image_size) |
| self._setup_indices() |
| self.subtask = "" |
| self.command = np.zeros(13, np.float32) |
| self._reset(randomize) |
| self._sample_command() |
| self.instruction = _instruction(spec, self.command, ordinal, self.subtask) |
| if instruction_override is not None: |
| self.instruction = instruction_override |
| self.autonomous_instruction = AUTONOMOUS_TASK_INSTRUCTIONS.get( |
| spec.name, self.instruction |
| ) |
| if self.started_fallen: |
| self.force_fall() |
|
|
| def _setup_indices(self) -> None: |
| import mujoco |
|
|
| m = self.model |
| self.trunk_body = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_BODY, "trunk_base") |
| self.trunk_joint = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_JOINT, "trunk_base_freejoint") |
| self.trunk_qpos = int(m.jnt_qposadr[self.trunk_joint]) |
| self.trunk_qvel = int(m.jnt_dofadr[self.trunk_joint]) |
| self.imu_sensor = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_SENSOR, "imu_ang_vel") |
| self.head_camera = mujoco.mj_name2id( |
| m, mujoco.mjtObj.mjOBJ_CAMERA, "head_camera" |
| ) |
| if self.head_camera < 0: |
| raise ValueError("scene does not contain head_camera") |
| m.cam_quat[self.head_camera] = HEAD_CAMERA_QUAT |
| self.joint_qpos = np.asarray( |
| [m.jnt_qposadr[m.actuator_trnid[i, 0]] for i in range(m.nu)], dtype=np.int64 |
| ) |
| self.joint_qvel = np.asarray( |
| [m.jnt_dofadr[m.actuator_trnid[i, 0]] for i in range(m.nu)], dtype=np.int64 |
| ) |
| if m.nu != ACTION_DIM: |
| raise ValueError(f"scene has {m.nu} actuators, expected {ACTION_DIM}") |
| ball_joint = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_JOINT, "ball_free") |
| self.ball_qpos = int(m.jnt_qposadr[ball_joint]) if ball_joint >= 0 else None |
| self.ball_qvel = int(m.jnt_dofadr[ball_joint]) if ball_joint >= 0 else None |
|
|
| def _reset(self, randomize: bool) -> None: |
| import mujoco |
|
|
| key_name = self.spec.start_keyframe |
| if self.spec.name == "sitstand" and self.ordinal % 2: |
| key_name = "SIT" |
| self.subtask = "stand_up" |
| elif self.spec.name == "sitstand": |
| self.subtask = "sit_then_stand" |
| key = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_KEY, key_name) |
| if key >= 0: |
| mujoco.mj_resetDataKeyframe(self.model, self.data, key) |
| else: |
| mujoco.mj_resetData(self.model, self.data) |
| self.data.qpos[self.trunk_qpos : self.trunk_qpos + 7] = [ |
| 0.0, |
| 0.0, |
| 0.1385 if "roller" in self.spec.name else 0.125, |
| 1.0, |
| 0.0, |
| 0.0, |
| 0.0, |
| ] |
| self.data.qpos[self.joint_qpos] = DEFAULT_POSE |
| self.data.ctrl[:] = DEFAULT_POSE |
| if randomize: |
| self.data.qpos[self.joint_qpos] += self.rng.uniform( |
| -math.radians(1.5), math.radians(1.5), ACTION_DIM |
| ) |
| yaw = self.rng.uniform(-math.radians(8), math.radians(8)) |
| self.data.qpos[self.trunk_qpos + 3 : self.trunk_qpos + 7] = [ |
| math.cos(yaw / 2), |
| 0.0, |
| 0.0, |
| math.sin(yaw / 2), |
| ] |
| self._randomize_model_visuals() |
| self.data.ctrl[:] = DEFAULT_POSE |
| self._set_wheel_friction() |
| mujoco.mj_forward(self.model, self.data) |
| if self.spec.kind == "kick": |
| self._place_ball() |
|
|
| def _randomize_model_visuals(self) -> None: |
| import mujoco |
|
|
| if self.model.nlight: |
| self.model.light_diffuse[:] *= self.rng.uniform(0.75, 1.2) |
| self.model.light_ambient[:] = np.clip( |
| self.model.light_ambient + self.rng.uniform(-0.08, 0.08), 0.0, 1.0 |
| ) |
| floor = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, "floor") |
| if floor >= 0: |
| rgb = self.rng.uniform(0.25, 0.75, 3) |
| self.model.geom_rgba[floor, :3] = rgb |
| if self.head_camera >= 0: |
| self.model.cam_pos[self.head_camera] += self.rng.uniform(-0.003, 0.003, 3) |
| self.model.cam_fovy[self.head_camera] = np.clip( |
| self.model.cam_fovy[self.head_camera] + self.rng.uniform(-3.0, 3.0), |
| 30.0, |
| 90.0, |
| ) |
|
|
| def _set_wheel_friction(self) -> None: |
| import mujoco |
|
|
| for j in range(self.model.njnt): |
| name = mujoco.mj_id2name(self.model, mujoco.mjtObj.mjOBJ_JOINT, j) |
| if name and name.startswith("passive_") and "wheel" in name: |
| self.model.dof_frictionloss[self.model.jnt_dofadr[j]] = 0.003 |
|
|
| def _sample_command(self) -> None: |
| if self.spec.name == "stand": |
| caps = np.array([0.35, 0.35, 0.45, 0.2], np.float32) |
| self.command[3:7] = self.rng.uniform(-caps, caps) |
| self.command[9] = self.rng.uniform(-0.015, 0.015) |
| self.command[10:12] = self.rng.uniform(-0.15, 0.15, 2) |
| elif self.spec.name == "walking": |
| candidates_x = np.array([-0.25, -0.15, 0.15, 0.25]) |
| candidates_y = np.array([-0.15, 0.0, 0.15]) |
| candidates_w = np.array([-1.0, -0.5, 0.0, 0.5, 1.0]) |
| self.command[:3] = [ |
| self.rng.choice(candidates_x), |
| self.rng.choice(candidates_y), |
| self.rng.choice(candidates_w), |
| ] |
| elif self.spec.name == "sitstand": |
| self.command[0] = 1.0 if self.subtask == "sit_then_stand" else 0.0 |
| elif self.spec.name == "roller": |
| self.command[0] = self.rng.uniform(0.25, 0.55) |
|
|
| def _place_ball(self) -> None: |
| if self.ball_qpos is None or self.ball_qvel is None: |
| raise ValueError("kick scene does not contain ball_free") |
| q = self.data.qpos[self.trunk_qpos + 3 : self.trunk_qpos + 7] |
| qw, qx, qy, qz = q |
| yaw = math.atan2(2 * (qw * qz + qx * qy), 1 - 2 * (qy * qy + qz * qz)) |
| off_y = ( |
| BALL_OFFSET_ABS_Y |
| if self.spec.name.endswith("left") |
| else -BALL_OFFSET_ABS_Y |
| ) |
| off_y += self.rng.uniform(-0.004, 0.004) |
| nominal_x = ( |
| BALL_OFFSET_X |
| if self.ball_forward_offset is None |
| else self.ball_forward_offset |
| ) |
| off_x = nominal_x + self.rng.uniform(-0.004, 0.004) |
| x, y = self.data.qpos[self.trunk_qpos : self.trunk_qpos + 2] |
| bx = x + math.cos(yaw) * off_x - math.sin(yaw) * off_y |
| by = y + math.sin(yaw) * off_x + math.cos(yaw) * off_y |
| self.data.qpos[self.ball_qpos : self.ball_qpos + 7] = [ |
| bx, |
| by, |
| BALL_RADIUS, |
| 1, |
| 0, |
| 0, |
| 0, |
| ] |
| self.data.qvel[self.ball_qvel : self.ball_qvel + 6] = 0 |
|
|
| def force_fall(self) -> None: |
| """Place the robot in a physically valid fallen pose for recovery data.""" |
| import mujoco |
|
|
| directions = ( |
| ("front", 1, 1.0), |
| ("back", 1, -1.0), |
| ("left", 0, 1.0), |
| ("right", 0, -1.0), |
| ) |
| name, axis, sign = directions[int(self.rng.integers(0, len(directions)))] |
| angle = sign * (math.pi / 2 + self.rng.uniform(-0.12, 0.12)) |
| quat = np.zeros(4, np.float64) |
| quat[0] = math.cos(angle / 2) |
| quat[1 + axis] = math.sin(angle / 2) |
| self.data.qpos[self.trunk_qpos + 2] = self.rng.uniform(0.048, 0.062) |
| self.data.qpos[self.trunk_qpos + 3 : self.trunk_qpos + 7] = quat |
| self.data.qpos[self.joint_qpos] = DEFAULT_POSE + self.rng.uniform( |
| -0.08, 0.08, ACTION_DIM |
| ) |
| self.data.qvel[:] = 0 |
| self.data.ctrl[:] = DEFAULT_POSE |
| mujoco.mj_forward(self.model, self.data) |
| self.fall_direction = name |
|
|
| def _observation(self, last_action: np.ndarray, command: np.ndarray) -> np.ndarray: |
| adr = self.model.sensor_adr[self.imu_sensor] |
| gyro = self.data.sensordata[adr : adr + 3].astype(np.float32).copy() |
| quat = self.data.xquat[self.trunk_body].astype(np.float32).copy() |
| gravity = _quat_rotate_inverse(quat, np.array([0, 0, -1], np.float32)) |
| joint_pos = self.data.qpos[self.joint_qpos].astype(np.float32) - DEFAULT_POSE |
| joint_vel = self.data.qvel[self.joint_qvel].astype(np.float32).copy() |
| obs = np.concatenate([gyro, gravity, joint_pos, joint_vel, last_action, command]) |
| if obs.shape != (STATE_DIM,) or not np.isfinite(obs).all(): |
| raise ValueError(f"invalid observation: shape={obs.shape}, finite={np.isfinite(obs).all()}") |
| return obs.astype(np.float32) |
|
|
| def _is_upright(self, gravity_z: float | None = None) -> bool: |
| if gravity_z is None: |
| quat = self.data.xquat[self.trunk_body].astype(np.float32).copy() |
| gravity_z = float( |
| _quat_rotate_inverse(quat, np.array([0, 0, -1], np.float32))[2] |
| ) |
| height = float(self.data.qpos[self.trunk_qpos + 2]) |
| return bool(gravity_z < -0.82 and height > 0.09) |
|
|
| def _recovery_detection_allowed(self, task_elapsed: float) -> bool: |
| if self.spec.kind not in ("kick", "roulade", "phase"): |
| return True |
| handoff = self.spec.handoff_after_s |
| return handoff is None or task_elapsed >= handoff |
|
|
| def _base_policy_filename_and_command(self, task_elapsed: float): |
| if self.control_program is not None: |
| return self.control_program(task_elapsed) |
| filename = self.spec.policy_file |
| command = self.command.copy() |
| if self.spec.name == "sitstand": |
| command[:] = 0 |
| if self.subtask == "sit_then_stand" and task_elapsed < 2.0: |
| command[0] = 1.0 |
| if self.spec.phase_period_s is not None and ( |
| self.spec.handoff_after_s is None or task_elapsed < self.spec.handoff_after_s |
| ): |
| phase = task_elapsed / self.spec.phase_period_s |
| command[:] = 0 |
| command[0] = math.cos(2 * math.pi * phase) |
| command[1] = math.sin(2 * math.pi * phase) |
| if ( |
| self.spec.handoff_after_s is not None |
| and task_elapsed >= self.spec.handoff_after_s |
| ): |
| filename = self.spec.handoff_policy_file or filename |
| command[:] = 0 |
| return filename, command, BALANCE_INSTRUCTION, "balance" |
| if self.spec.kind == "recovery": |
| return "alpha_stand.onnx", np.zeros_like(command), BALANCE_INSTRUCTION, "balance" |
| return filename, command, self.instruction, "task" |
|
|
| def _control_context(self, step: int): |
| elapsed = step / CONTROL_HZ |
| task_elapsed = ( |
| 0.0 |
| if self.task_start_step is None |
| else max(0.0, (step - self.task_start_step) / CONTROL_HZ) |
| ) |
| quat = self.data.xquat[self.trunk_body].astype(np.float32).copy() |
| gravity_z = float( |
| _quat_rotate_inverse(quat, np.array([0, 0, -1], np.float32))[2] |
| ) |
| height = float(self.data.qpos[self.trunk_qpos + 2]) |
| upright = self._is_upright(gravity_z) |
|
|
| if self.recovery_active: |
| self.recovery_upright_ticks = ( |
| self.recovery_upright_ticks + 1 if upright else 0 |
| ) |
| required = round(RECOVERY_STABLE_S * CONTROL_HZ) |
| if self.recovery_upright_ticks >= required: |
| self.recovery_active = False |
| self.recovery_completed = True |
| self.recovery_exit_step = step |
| self.task_start_step = step |
| self.recovery_upright_ticks = 0 |
| task_elapsed = 0.0 |
| elif ( |
| elapsed < self.spec.duration_s - TERMINAL_STABILIZATION_S |
| and self._recovery_detection_allowed(task_elapsed) |
| and (gravity_z > -0.58 or height < 0.073) |
| ): |
| self.recovery_active = True |
| self.recovery_enter_step = step |
| self.recovery_upright_ticks = 0 |
|
|
| if self.recovery_active: |
| return ( |
| "alpha_stand.onnx", |
| np.zeros_like(self.command), |
| RECOVERY_INSTRUCTION, |
| "recovery", |
| ) |
| if elapsed >= self.spec.duration_s - TERMINAL_STABILIZATION_S: |
| return ( |
| "alpha_stand.onnx", |
| np.zeros_like(self.command), |
| BALANCE_INSTRUCTION, |
| "balance", |
| ) |
| context = self._base_policy_filename_and_command(task_elapsed) |
| if context[3] == "task": |
| self.task_control_steps += 1 |
| return context |
|
|
| def run(self, delay_ticks: int = 0) -> EpisodeResult: |
| import mujoco |
|
|
| total_steps = round(self.spec.duration_s * CONTROL_HZ) |
| last_action = np.zeros(ACTION_DIM, np.float32) |
| delay = [np.zeros(ACTION_DIM, np.float32) for _ in range(delay_ticks + 1)] |
| images: list[np.ndarray] = [] |
| states: list[np.ndarray] = [] |
| obs_indices: list[int] = [] |
| observation_instructions: list[str] = [] |
| observation_phases: list[str] = [] |
| actions: list[np.ndarray] = [] |
| trunk_xyz: list[np.ndarray] = [] |
| gravity_z: list[float] = [] |
| cumulative_abs_pitch_rate = 0.0 |
| ball_start = ( |
| self.data.qpos[self.ball_qpos : self.ball_qpos + 3].copy() |
| if self.ball_qpos is not None |
| else None |
| ) |
| previous_program_policy: str | None = None |
|
|
| for step in range(total_steps): |
| filename, command, active_instruction, active_phase = self._control_context(step) |
| if ( |
| self.control_program is not None |
| and previous_program_policy is not None |
| and filename != previous_program_policy |
| ): |
| |
| |
| |
| |
| last_action = np.zeros(ACTION_DIM, np.float32) |
| previous_program_policy = filename |
| session = self.sessions.get(filename) |
| obs = self._observation(last_action, command) |
| inp = session.get_inputs()[0].name |
| out = session.get_outputs()[0].name |
| action = session.run([out], {inp: obs[None]})[0][0].astype(np.float32) |
| if action.shape != (ACTION_DIM,) or not np.isfinite(action).all(): |
| raise ValueError(f"invalid action from {self.spec.name}: {action}") |
| if step % CONTROL_DECIMATION == 0: |
| self.renderer.update_scene(self.data, camera="head_camera") |
| image = self.renderer.render().copy() |
| if image.shape != (self.image_size, self.image_size, 3): |
| raise ValueError(f"unexpected render shape {image.shape}") |
| images.append(image) |
| states.append(obs) |
| obs_indices.append(step) |
| observation_instructions.append(active_instruction) |
| observation_phases.append(active_phase) |
| actions.append(action) |
| delay.append(action.copy()) |
| applied = delay.pop(0) |
| self.data.ctrl[:] = DEFAULT_POSE + applied |
| for _ in range(4): |
| mujoco.mj_step(self.model, self.data) |
| last_action = action |
| trunk_xyz.append(self.data.qpos[self.trunk_qpos : self.trunk_qpos + 3].copy()) |
| gravity_z.append(float(self._observation(last_action, command)[5])) |
| cumulative_abs_pitch_rate += abs(float(self.data.qvel[self.trunk_qvel + 4])) / CONTROL_HZ |
|
|
| return self._finish_result( |
| images, |
| states, |
| obs_indices, |
| actions, |
| trunk_xyz, |
| gravity_z, |
| cumulative_abs_pitch_rate, |
| ball_start, |
| observation_instructions, |
| observation_phases, |
| None, |
| ) |
|
|
| def run_chunk_policy( |
| self, |
| policy, |
| delay_ticks: int = 0, |
| record_external: bool = False, |
| implicit_phase_instruction: bool = False, |
| ) -> EpisodeResult: |
| """Run a vision policy that replans at the 10 Hz observation rate. |
| |
| ``policy(image, state, instruction, control_index)`` returns at least |
| five raw 14D joint-offset actions. The simulator executes the first |
| five at 50 Hz, then captures a fresh head-camera observation. |
| """ |
| import mujoco |
|
|
| total_steps = round(self.spec.duration_s * CONTROL_HZ) |
| last_action = np.zeros(ACTION_DIM, np.float32) |
| delay = [np.zeros(ACTION_DIM, np.float32) for _ in range(delay_ticks + 1)] |
| images: list[np.ndarray] = [] |
| states: list[np.ndarray] = [] |
| obs_indices: list[int] = [] |
| observation_instructions: list[str] = [] |
| observation_phases: list[str] = [] |
| video_frames: list[np.ndarray] = [] |
| actions: list[np.ndarray] = [] |
| trunk_xyz: list[np.ndarray] = [] |
| gravity_z: list[float] = [] |
| cumulative_abs_pitch_rate = 0.0 |
| ball_start = ( |
| self.data.qpos[self.ball_qpos : self.ball_qpos + 3].copy() |
| if self.ball_qpos is not None |
| else None |
| ) |
| active_chunk: np.ndarray | None = None |
| tracking_camera = None |
| if record_external: |
| tracking_camera = mujoco.MjvCamera() |
| mujoco.mjv_defaultCamera(tracking_camera) |
| tracking_camera.type = mujoco.mjtCamera.mjCAMERA_TRACKING |
| tracking_camera.trackbodyid = self.trunk_body |
| tracking_camera.distance = 0.62 |
| tracking_camera.azimuth = 135.0 |
| tracking_camera.elevation = -18.0 |
|
|
| for step in range(total_steps): |
| _, command, active_instruction, active_phase = self._control_context(step) |
| policy_instruction = ( |
| self.autonomous_instruction |
| if implicit_phase_instruction |
| else active_instruction |
| ) |
| obs = self._observation(last_action, command) |
| if step % CONTROL_DECIMATION == 0: |
| self.renderer.update_scene(self.data, camera="head_camera") |
| image = self.renderer.render().copy() |
| active_chunk = np.asarray( |
| policy(image, obs, policy_instruction, step), dtype=np.float32 |
| ) |
| if ( |
| active_chunk.ndim != 2 |
| or active_chunk.shape[0] < CONTROL_DECIMATION |
| or active_chunk.shape[1] != ACTION_DIM |
| or not np.isfinite(active_chunk).all() |
| ): |
| raise ValueError(f"vision policy returned invalid chunk {active_chunk.shape}") |
| images.append(image) |
| states.append(obs) |
| obs_indices.append(step) |
| observation_instructions.append(policy_instruction) |
| observation_phases.append(active_phase) |
| if tracking_camera is not None: |
| self.renderer.update_scene(self.data, camera=tracking_camera) |
| video_frames.append(self.renderer.render().copy()) |
| assert active_chunk is not None |
| action = active_chunk[step % CONTROL_DECIMATION] |
| actions.append(action.copy()) |
| delay.append(action.copy()) |
| applied = delay.pop(0) |
| self.data.ctrl[:] = DEFAULT_POSE + applied |
| for _ in range(4): |
| mujoco.mj_step(self.model, self.data) |
| last_action = action |
| trunk_xyz.append(self.data.qpos[self.trunk_qpos : self.trunk_qpos + 3].copy()) |
| gravity_z.append(float(self._observation(last_action, command)[5])) |
| cumulative_abs_pitch_rate += ( |
| abs(float(self.data.qvel[self.trunk_qvel + 4])) / CONTROL_HZ |
| ) |
|
|
| return self._finish_result( |
| images, |
| states, |
| obs_indices, |
| actions, |
| trunk_xyz, |
| gravity_z, |
| cumulative_abs_pitch_rate, |
| ball_start, |
| observation_instructions, |
| observation_phases, |
| np.asarray(video_frames, np.uint8) if video_frames else None, |
| ) |
|
|
| def _finish_result( |
| self, |
| images, |
| states, |
| obs_indices, |
| actions, |
| trunk_xyz, |
| gravity_z, |
| cumulative_abs_pitch_rate, |
| ball_start, |
| observation_instructions, |
| observation_phases, |
| video_frames, |
| ) -> EpisodeResult: |
| self.renderer.close() |
| xyz = np.asarray(trunk_xyz) |
| upright = (np.asarray(gravity_z) < -0.72) & (xyz[:, 2] > 0.09) |
| terminal_steps = max(1, round(TERMINAL_STABILIZATION_S * CONTROL_HZ / 2)) |
| terminal_upright_fraction = float(upright[-terminal_steps:].mean()) |
| terminal_stable = bool(terminal_upright_fraction >= 0.8 and upright[-1]) |
| final_upright = bool(upright[-1]) |
| planar_displacement = float(np.linalg.norm(xyz[-1, :2] - xyz[0, :2])) |
| planar_path_length = float( |
| np.linalg.norm(np.diff(xyz[:, :2], axis=0), axis=1).sum() |
| ) |
| min_height = float(xyz[:, 2].min()) |
| ball_displacement = 0.0 |
| if ball_start is not None: |
| ball_end = self.data.qpos[self.ball_qpos : self.ball_qpos + 3] |
| ball_displacement = float(np.linalg.norm(ball_end[:2] - ball_start[:2])) |
|
|
| if self.spec.kind == "recovery": |
| success = bool(self.recovery_completed and terminal_stable) |
| elif self.started_fallen: |
| success = bool( |
| self.recovery_completed |
| and self.task_control_steps >= CONTROL_HZ // 2 |
| and terminal_stable |
| ) |
| elif self.spec.name == "sitstand" and self.subtask == "sit_then_stand": |
| success = bool(min_height < xyz[0, 2] - 0.008 and terminal_stable) |
| elif self.spec.kind == "kick": |
| success = bool(ball_displacement > 0.03 and terminal_stable) |
| elif self.spec.kind == "roulade": |
| success = bool(cumulative_abs_pitch_rate > 2.5 and terminal_stable) |
| elif self.spec.name in ("walking", "roller"): |
| success = bool(planar_displacement > 0.025 and terminal_stable) |
| elif self.spec.kind == "phase": |
| success = bool(min_height < xyz[0, 2] - 0.01 and terminal_stable) |
| else: |
| success = bool(terminal_stable) |
|
|
| metrics: dict[str, float | bool | str] = { |
| "success": success, |
| "final_upright": final_upright, |
| "planar_displacement_m": planar_displacement, |
| "planar_path_length_m": planar_path_length, |
| "initial_trunk_height_m": float(xyz[0, 2]), |
| "minimum_trunk_height_m": min_height, |
| "final_trunk_height_m": float(xyz[-1, 2]), |
| "ball_displacement_m": ball_displacement, |
| "integrated_abs_pitch_rate_rad": cumulative_abs_pitch_rate, |
| "subtask": self.subtask, |
| "terminal_upright_fraction": terminal_upright_fraction, |
| "terminal_stable": terminal_stable, |
| "started_fallen": self.started_fallen, |
| "fall_direction": self.fall_direction, |
| "recovery_completed": self.recovery_completed, |
| "recovery_time_s": ( |
| None |
| if self.recovery_exit_step is None |
| else self.recovery_exit_step / CONTROL_HZ |
| ), |
| "task_control_steps": self.task_control_steps, |
| } |
| return EpisodeResult( |
| images=np.asarray(images, np.uint8), |
| states=np.asarray(states, np.float32), |
| observation_control_indices=np.asarray(obs_indices, np.int32), |
| actions=np.asarray(actions, np.float32), |
| instruction=self.instruction, |
| success=success, |
| metrics=metrics, |
| command=self.command, |
| subtask=self.subtask, |
| observation_instructions=tuple(observation_instructions), |
| observation_phases=tuple(observation_phases), |
| video_frames=video_frames, |
| ) |
|
|
|
|
| def encode_episode_metadata(metadata: dict) -> np.ndarray: |
| def convert(value): |
| if isinstance(value, np.generic): |
| return value.item() |
| if isinstance(value, np.ndarray): |
| return value.tolist() |
| raise TypeError(f"cannot JSON-encode {type(value).__name__}") |
|
|
| return np.asarray(json.dumps(metadata, ensure_ascii=False, default=convert)) |
|
|
|
|
| def decode_episode_metadata(value: np.ndarray) -> dict: |
| return json.loads(str(value.item())) |
|
|
