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 # The generated MJCF camera quaternion points the optical axis into the head # shell. MuJoCo cameras look down local -Z, so use the physical convention: # robot +X is forward, world +Z is image-up, with 15 degrees of downward tilt. 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 ): # The source ONNX teachers are stateless policies trained from # their own neutral previous-action convention. Reset only the # observation feature at a composed teacher handoff; physical # controls still pass continuously through the delay buffer. 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()))