Hugging Face's logo

MtnMCG
/
PCEI-Robotics

Robotics
English
embodied-ai
mujoco
humanoid
reinforcement-learning
evolutionary-strategies
pcei
mamba
bitnet
fastkan
Model card Files
xet
 Community
PCEI-Robotics / ninja_env.py
MtnMCG's picture
MtnMCG
Upload ninja_env.py with huggingface_hub
a1d2967 verified
Raw History Blame Contribute Delete
94.4 kB
"""
embodied_pcei/src/ninja_env.py
High-performance physical simulation environment for the American Ninja Warrior course.
Built on MuJoCo with native Apple Silicon performance.
Features:
- Articulated bipedal humanoid with 21 joint torque actuators.
- Staged American Ninja Warrior course:
* Stage 0: Flat sprint runway (0 to 8m)
* Stage 1: Quad steps / angled stepping stones (8 to 18m)
* Stage 2: Elevated balance beam (18 to 28m)
* Stage 3: Hurdles & stepping barriers (28 to 38m)
* Stage 4: Warped wall incline ramp & summit buzzer (38 to 48m)
- Raycast Pseudo-Vision: forward height probes measuring terrain/obstacle topology.
- Metabolic / Homeostatic Telemetry: lungs/stamina, stomach/hunger, heart/health.
- Built-in offscreen renderer for live web dashboard streaming.
"""
import numpy as np
import mujoco
from pathlib import Path
from typing import Dict, Tuple, Optional
XML_PATH = Path(__file__).resolve().parent.parent / "models" / "ninja_course.xml"
FLAT_ARENA_XML = Path(__file__).resolve().parent.parent / "models" / "flat_arena.xml"
# Obstacle Stage Milestones along +X
STAGE_MILESTONES = [
(8.0, "Stage 1: Quad Steps"),
(18.0, "Stage 2: Balance Beam"),
(28.0, "Stage 3: Hurdles"),
(38.0, "Stage 4: Warped Wall"),
(46.0, "Course Finish: Summit Buzzer"),
]
# 18 Progressive Navigation Waypoints across the 48m Ninja Warrior Course
# Format: (Target 3D Coordinate [X, Y, Z], Waypoint Label, Arrival Radius in meters)
COURSE_WAYPOINTS = [
(np.array([3.0, 0.00, 1.33], dtype=np.float32), "Runway Accelerate 1", 1.2),
(np.array([6.0, 0.00, 1.33], dtype=np.float32), "Runway Accelerate 2", 1.2),
(np.array([8.0, 0.00, 1.33], dtype=np.float32), "Runway Exit Threshold", 1.0),
(np.array([10.0, -0.60, 0.20], dtype=np.float32), "Stage 1: Quad Step 1 (Right)", 0.8),
(np.array([12.2, 0.60, 0.30], dtype=np.float32), "Stage 1: Quad Step 2 (Left)", 0.8),
(np.array([14.4, -0.50, 0.35], dtype=np.float32), "Stage 1: Quad Step 3 (Right)", 0.8),
(np.array([16.6, 0.50, 0.40], dtype=np.float32), "Stage 1: Quad Step 4 (Left)", 0.8),
(np.array([18.5, 0.00, 0.45], dtype=np.float32), "Platform Mid 1", 1.0),
(np.array([20.5, 0.00, 0.50], dtype=np.float32), "Stage 2: Balance Beam Entry", 0.6),
(np.array([23.5, 0.00, 0.50], dtype=np.float32), "Stage 2: Balance Beam Mid", 0.6),
(np.array([26.5, 0.00, 0.50], dtype=np.float32), "Stage 2: Balance Beam Exit", 0.6),
(np.array([28.5, 0.00, 0.50], dtype=np.float32), "Platform Mid 2", 1.0),
(np.array([31.5, 0.00, 0.20], dtype=np.float32), "Stage 3: Hurdle 1", 0.9),
(np.array([34.5, 0.00, 0.28], dtype=np.float32), "Stage 3: Hurdle 2", 0.9),
(np.array([37.5, 0.00, 0.22], dtype=np.float32), "Stage 3: Hurdle 3", 0.9),
(np.array([40.0, 0.00, 0.30], dtype=np.float32), "Stage 4: Warped Wall Approach",1.2),
(np.array([42.5, 0.00, 0.70], dtype=np.float32), "Stage 4: Warped Wall Incline", 1.2),
(np.array([47.5, 0.00, 2.45], dtype=np.float32), "Course Finish: Summit Buzzer", 0.6),
]
# Translucent Neon Color Palette for Multi-Agent Ghost Swarm Rendering
GHOST_PALETTE = [
np.array([0.2, 0.7, 1.0, 0.16], dtype=np.float32), # Neon Cyan
np.array([0.7, 0.3, 1.0, 0.16], dtype=np.float32), # Electric Purple
np.array([0.1, 0.9, 0.6, 0.16], dtype=np.float32), # Neon Emerald
np.array([1.0, 0.55, 0.1, 0.16], dtype=np.float32), # Sunset Orange
np.array([1.0, 0.2, 0.6, 0.16], dtype=np.float32), # Neon Pink
np.array([0.35, 0.8, 1.0, 0.16], dtype=np.float32), # Sky Blue
np.array([0.85, 0.95, 0.2, 0.16], dtype=np.float32),# Lime Volt
]
# Canonical Reference Quiet Standing Equilibrium Pose (21 Actuated Joints)
# Torso at Z=1.168m, spine vertical, knees naturally flexed at 0.05 rad (~3 deg), arms down at sides
STAND_REF_QPOS = np.array([
0.0, 0.0, 0.0, # abdomen_z, abdomen_y, abdomen_x
0.0, 0.0, -0.05, 0.05, -0.025, 0.0, # right_hip_x, z, y, knee, ankle_y, ankle_x
0.0, 0.0, -0.05, 0.05, -0.025, 0.0, # left_hip_x, z, y, knee, ankle_y, ankle_x
0.65, 0.0, 0.0, # right_shoulder_x, y, elbow (down at side)
-0.65, 0.0, 0.0, # left_shoulder_x, y, elbow (down at side)
], dtype=np.float32)
STAND_JOINT_WEIGHTS = np.array([
1.5, 1.5, 1.5, # abdomen_z, abdomen_y, abdomen_x
1.5, 1.5, 1.5, 1.5, 1.5, 1.5, # right leg (hip_x, hip_z, hip_y, knee, ankle_y, ankle_x)
1.5, 1.5, 1.5, 1.5, 1.5, 1.5, # left leg (hip_x, hip_z, hip_y, knee, ankle_y, ankle_x)
0.5, 0.5, 0.5, # right arm (shoulder_x, shoulder_y, elbow)
0.5, 0.5, 0.5 # left arm (shoulder_x, shoulder_y, elbow)
], dtype=np.float32)
class NinjaEnv:
def __init__(self, xml_path: Optional[str] = None, frame_skip: int = 5, arena_mode: bool = True, stream_mode: str = "cruise"):
self.arena_mode = arena_mode
self.stream_mode = stream_mode # "cruise", "stand", "getup", "reach"
if xml_path is None:
xml_path = str(FLAT_ARENA_XML if arena_mode else XML_PATH)
self.model = mujoco.MjModel.from_xml_path(xml_path)
self.data = mujoco.MjData(self.model)
self.frame_skip = frame_skip
self.dt = self.model.opt.timestep * frame_skip # 0.005 * 5 = 0.025s (40 Hz)
self.n_actuators = self.model.nu # 21
self.n_qpos = self.model.nq # 28 (7 root + 21 joints)
self.n_qvel = self.model.nv # 27 (6 root + 21 joints)
# Solid-State Angular LiDAR Sensor Configuration
# 16 vertical beams (pitch: -35 deg to +15 deg) x 32 horizontal beams (yaw: -45 deg to +45 deg) = 512 rays
self.num_lidar_v = 16
self.num_lidar_h = 32
self.num_lidar_rays = self.num_lidar_v * self.num_lidar_h # 512
self.num_rays = self.num_lidar_rays # backwards-compatibility alias
self.lidar_max_range = 10.0 # meters
self.lidar_v_angles = np.radians(np.linspace(-35, 15, self.num_lidar_v))
self.lidar_h_angles = np.radians(np.linspace(-45, 45, self.num_lidar_h))
# Precompute unit direction vectors in head local coordinate frame
# Forward = +X, Left = +Y, Up = +Z
local_vectors = []
for pitch in self.lidar_v_angles:
for yaw in self.lidar_h_angles:
vx = np.cos(pitch) * np.cos(yaw)
vy = np.cos(pitch) * np.sin(yaw)
vz = np.sin(pitch)
v = np.array([vx, vy, vz], dtype=np.float64)
v /= np.linalg.norm(v)
local_vectors.append(v)
self.lidar_local_vectors = np.array(local_vectors, dtype=np.float64) # (128, 3)
self.current_lidar_grid = np.zeros((self.num_lidar_v, self.num_lidar_h), dtype=np.float32)
self._geomid_buf = np.zeros(1, dtype=np.int32)
self._visor_offset = np.array([0.15, 0.0, 0.0], dtype=np.float64)
# Metabolic internal states
self.stamina = 1.0 # 0.0 to 1.0
self.health = 1.0 # 0.0 to 1.0
self.hunger = 0.0 # 0.0 to 1.0
# Milestone tracking & Course Waypoints
self.milestones_reached = set()
self.waypoints_reached = set()
self.current_wp_idx = 0
self.furthest_x = 0.0
self.max_steps = 1000
# Flat Arena & Dynamic Random Goal Navigation State (Stream 4: Cruise)
self.goal_pos = np.array([5.0, 0.0, 1.25], dtype=np.float32)
self.goals_completed = 0
self.goal_arrival_radius = 2.0 # Wide capture radius (2.0m): easily clears by walking near target
self.min_dist_to_goal = 5.0
self.prev_dist_to_goal = 5.0
self.steps_not_approaching = 0
self.transition_leeway_window = 90 # ~2.25s leeway (90 steps at 40Hz) for turning & realignment
self.goal_transition_steps = 0 # No leeway on initial spawn (must advance immediately)
# Stream 1: Quiet Standing Stillness State
self.stand_still_steps = 0
self.stand_drift_steps = 0
self.stand_one_leg_steps = 0
# Stream 2: Fall Recovery State (Ragdoll & Curated Library)
self.fall_state_type = "random_ragdoll"
self.consecutive_upright_steps = 0
self.prev_z = 0.20
self.prev_head_z = 0.20
self.getup_stagnant_steps = 0
self.max_elevation_achieved = 0.20
self.max_head_elevation = 0.20
self.getup_goal_pos = np.array([0.0, 0.0, 1.168], dtype=np.float32)
self.getup_success_count = 0
# Stream 3: Bilateral Hand Reaching State (Bolted Torso)
self.active_reach_hand = "right"
self.reach_target_pos = np.array([0.35, -0.25, 1.35], dtype=np.float32)
self.reach_touches = 0
self.prev_reach_dist = 0.40
self.min_reach_dist = 0.40
self.steps_on_reach_target = 0
self.steps_receding_reach = 0
# Bilateral gait stance duty-cycle counters
self.r_foot_contact_steps = 0
self.l_foot_contact_steps = 0
# Offscreen renderer instance (lazy-loaded)
self._renderer = None
# Cache body and geom IDs for fast lookups
self.torso_body_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_BODY, "torso")
self.head_body_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_BODY, "head")
self.r_foot_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_BODY, "right_foot")
self.l_foot_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_BODY, "left_foot")
self.r_foot_geom_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, "right_foot_geom")
self.l_foot_geom_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, "left_foot_geom")
self.r_hand_geom_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, "right_hand")
self.l_hand_geom_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, "left_hand")
self.r_arm_body_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_BODY, "right_lower_arm")
self.l_arm_body_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_BODY, "left_lower_arm")
# Cache leg joint qpos indices for athletic stance and gait monitoring
self.r_knee_qpos = self.model.jnt_qposadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "right_knee")]
self.l_knee_qpos = self.model.jnt_qposadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "left_knee")]
self.r_hip_qpos = self.model.jnt_qposadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "right_hip_y")]
self.l_hip_qpos = self.model.jnt_qposadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "left_hip_y")]
self.r_ankle_qpos = self.model.jnt_qposadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "right_ankle_y")]
self.l_ankle_qpos = self.model.jnt_qposadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "left_ankle_y")]
# Cache abdomen joint qpos indices for core posture alignment
self.abdomen_x_qpos = self.model.jnt_qposadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "abdomen_x")]
self.abdomen_y_qpos = self.model.jnt_qposadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "abdomen_y")]
self.abdomen_z_qpos = self.model.jnt_qposadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "abdomen_z")]
# Cache arm joint qpos and dof velocity indices for bilateral reach
self.arm_qpos_indices = [
self.model.jnt_qposadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "right_shoulder_x")],
self.model.jnt_qposadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "right_shoulder_y")],
self.model.jnt_qposadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "right_elbow")],
self.model.jnt_qposadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "left_shoulder_x")],
self.model.jnt_qposadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "left_shoulder_y")],
self.model.jnt_qposadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "left_elbow")],
]
self.arm_dof_indices = [
self.model.jnt_dofadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "right_shoulder_x")],
self.model.jnt_dofadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "right_shoulder_y")],
self.model.jnt_dofadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "right_elbow")],
self.model.jnt_dofadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "left_shoulder_x")],
self.model.jnt_dofadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "left_shoulder_y")],
self.model.jnt_dofadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "left_elbow")],
]
# Cache collision geoms for physical impact detection (head, visor, chest, torso)
# All geoms belonging to world body (body 0) are environmental surfaces: floor + all obstacles
self.world_obstacle_geoms = {i for i in range(self.model.ngeom) if self.model.geom_bodyid[i] == 0}
self.body_impact_geoms = {
mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, "torso_geom"),
mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, "chest_geom"),
mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, "head_geom"),
mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, "visor"),
}
# Cache humanoid geom indices (all geoms belonging to non-world bodies) for ghost rendering
self.humanoid_geoms = [i for i in range(self.model.ngeom) if self.model.geom_bodyid[i] != 0]
# Cache limb geoms for Get-Up multi-phase floor contact righting
self.r_shin_geom_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, "right_shin_geom")
self.l_shin_geom_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, "left_shin_geom")
self.support_hand_geoms = {
self.r_hand_geom_id,
self.l_hand_geom_id,
mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, "right_larm_geom"),
mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, "left_larm_geom"),
}
self.support_knee_geoms = {self.r_shin_geom_id, self.l_shin_geom_id}
def _lock_torso_for_reach(self):
"""Bolts root torso, core spine, and legs in neutral athletic stance (Stream 3: Isolated Arm Motor Cortex)."""
self.data.qpos[0] = 0.0
self.data.qpos[1] = 0.0
self.data.qpos[2] = 1.35
self.data.qpos[3:7] = np.array([1.0, 0.0, 0.0, 0.0]) # Upright quaternion
self.data.qvel[0:6] = 0.0 # Zero base linear & angular velocity
# Lock core spine neutral
self.data.qpos[7:10] = 0.0 # abdomen_z, abdomen_y, abdomen_x
# Lock legs in neutral supportive stance
self.data.qpos[10:22] = 0.0 # zero all hip, knee, ankle
self.data.qpos[self.r_knee_qpos] = 0.30
self.data.qpos[self.l_knee_qpos] = 0.30
self.data.qpos[self.r_ankle_qpos] = -0.10
self.data.qpos[self.l_ankle_qpos] = -0.10
self.data.qvel[6:21] = 0.0 # Zero core & lower body velocities (DOFs 6 to 20, preserving arm DOFs 21:27)
def _sample_reach_target(self):
"""
Samples a randomized 3D reach goal via Forward Kinematics (FK) from the active arm's
reachable joint space, mathematically guaranteeing 100.0% of targets are physically reachable (Stream 3).
"""
self.active_reach_hand = np.random.choice(["right", "left"])
# Save current posture
saved_qpos = self.data.qpos.copy()
target_pos = None
# Sample joint angles from the arm's valid forward/lateral reachable envelope
for _ in range(25):
if self.active_reach_hand == "right":
# Right arm envelope: shoulder roll/pitch and elbow flexion
target_sx = float(np.random.uniform(-0.75, 0.40))
target_sy = float(np.random.uniform(-0.60, 0.85))
target_el = float(np.random.uniform(-1.10, 0.45))
self.data.qpos[self.arm_qpos_indices[0]] = target_sx # right_shoulder_x
self.data.qpos[self.arm_qpos_indices[1]] = target_sy # right_shoulder_y
self.data.qpos[self.arm_qpos_indices[2]] = target_el # right_elbow
mujoco.mj_forward(self.model, self.data)
cand_pos = self.data.geom_xpos[self.r_hand_geom_id].copy()
else:
# Left arm envelope: shoulder roll/pitch and elbow flexion
target_sx = float(np.random.uniform(-0.40, 0.75))
target_sy = float(np.random.uniform(-0.60, 0.85))
target_el = float(np.random.uniform(-0.45, 1.10))
self.data.qpos[self.arm_qpos_indices[3]] = target_sx # left_shoulder_x
self.data.qpos[self.arm_qpos_indices[4]] = target_sy # left_shoulder_y
self.data.qpos[self.arm_qpos_indices[5]] = target_el # left_elbow
mujoco.mj_forward(self.model, self.data)
cand_pos = self.data.geom_xpos[self.l_hand_geom_id].copy()
# Restore original pose to measure distance
self.data.qpos[:] = saved_qpos
mujoco.mj_forward(self.model, self.data)
curr_h = self.data.geom_xpos[self.r_hand_geom_id] if self.active_reach_hand == "right" else self.data.geom_xpos[self.l_hand_geom_id]
dist_to_cand = float(np.linalg.norm(curr_h - cand_pos))
target_pos = cand_pos
# Ensure target is at least 0.16m away so the agent must actively articulate its arm
if dist_to_cand >= 0.16:
break
# Always restore state cleanly
self.data.qpos[:] = saved_qpos
mujoco.mj_forward(self.model, self.data)
self.reach_target_pos = target_pos.astype(np.float32)
self.goal_pos = self.reach_target_pos.copy()
def _apply_fall_state(self, pose_name: Optional[str] = None):
"""
Initializes agent in a randomized fall state with 360Β° randomized yaw,
perturbed joints, physics contact settling, and a randomized standing destination goal (Stream 2).
"""
if pose_name is None:
poses = ["prone", "supine", "lateral_left", "lateral_right", "kneeling_crumpled", "seated_slumped", "random_ragdoll"]
pose_name = np.random.choice(poses)
self.fall_state_type = pose_name
self.data.qvel[:] = 0.0
# Random initial floor translation offset
spawn_x = float(np.random.uniform(-0.4, 0.4))
spawn_y = float(np.random.uniform(-0.4, 0.4))
# Base orientation quaternion [w, x, y, z] per canonical pose
if pose_name == "prone":
# Lying on stomach: pitch 90 deg forward
self.data.qpos[0:3] = np.array([spawn_x, spawn_y, 0.22])
q_base = np.array([0.7071, 0.0, 0.7071, 0.0])
self.data.qpos[self.r_knee_qpos] = 0.2
self.data.qpos[self.l_knee_qpos] = 0.2
elif pose_name == "supine":
# Lying on back: pitch 90 deg backward
self.data.qpos[0:3] = np.array([spawn_x, spawn_y, 0.22])
q_base = np.array([0.7071, 0.0, -0.7071, 0.0])
self.data.qpos[self.r_knee_qpos] = 0.3
self.data.qpos[self.l_knee_qpos] = 0.3
elif pose_name == "lateral_left":
# Lying on left flank: roll 90 deg left
self.data.qpos[0:3] = np.array([spawn_x, spawn_y, 0.22])
q_base = np.array([0.7071, -0.7071, 0.0, 0.0])
elif pose_name == "lateral_right":
# Lying on right flank: roll 90 deg right
self.data.qpos[0:3] = np.array([spawn_x, spawn_y, 0.22])
q_base = np.array([0.7071, 0.7071, 0.0, 0.0])
elif pose_name == "kneeling_crumpled":
# Crumpled on knees, torso pitched forward
self.data.qpos[0:3] = np.array([spawn_x, spawn_y, 0.48])
q_base = np.array([0.9238, 0.0, 0.3826, 0.0])
self.data.qpos[self.r_knee_qpos] = 1.7
self.data.qpos[self.l_knee_qpos] = 1.7
self.data.qpos[self.r_hip_qpos] = -1.0
self.data.qpos[self.l_hip_qpos] = -1.0
elif pose_name == "seated_slumped":
# Slumped seated on buttocks, legs forward
self.data.qpos[0:3] = np.array([spawn_x, spawn_y, 0.35])
q_base = np.array([0.9238, 0.0, -0.3826, 0.0])
self.data.qpos[self.r_hip_qpos] = -1.4
self.data.qpos[self.l_hip_qpos] = -1.4
self.data.qpos[self.r_knee_qpos] = 0.2
self.data.qpos[self.l_knee_qpos] = 0.2
else:
# Random ragdoll drop: drop from 0.45m with completely randomized 3D orientation
self.data.qpos[0:3] = np.array([spawn_x, spawn_y, float(np.random.uniform(0.40, 0.60))])
rand_quat = np.random.uniform(-1.0, 1.0, size=4)
q_base = rand_quat / np.linalg.norm(rand_quat)
# Apply randomized 360Β° yaw rotation to base pose quaternion
yaw = float(np.random.uniform(-np.pi, np.pi))
q_yaw = np.array([np.cos(yaw / 2.0), 0.0, 0.0, np.sin(yaw / 2.0)], dtype=np.float64)
# Quaternion multiplication: q_yaw * q_base
w1, x1, y1, z1 = q_yaw
w2, x2, y2, z2 = q_base
q_rot = np.array([
w1*w2 - x1*x2 - y1*y2 - z1*z2,
w1*x2 + x1*w2 + y1*z2 - z1*y2,
w1*y2 - x1*z2 + y1*w2 + z1*x2,
w1*z2 + x1*y2 - y1*x2 + z1*w2,
], dtype=np.float64)
self.data.qpos[3:7] = q_rot / np.linalg.norm(q_rot)
# Stochastic joint perturbation across all 21 joint degrees of freedom
self.data.qpos[7:] += np.random.uniform(-0.20, 0.20, size=self.n_actuators)
# Settle contact forces smoothly (12 passive sub-steps with zero actuation)
mujoco.mj_forward(self.model, self.data)
for _ in range(12):
self.data.ctrl[:] = 0.0
mujoco.mj_step(self.model, self.data)
self.data.qvel[:] = 0.0
mujoco.mj_forward(self.model, self.data)
# Sample a randomized standing destination goal beacon (0.6m - 1.8m away at random heading)
stand_dist = float(np.random.uniform(0.6, 1.8))
stand_heading = float(np.random.uniform(-np.pi, np.pi))
cur_torso = self.data.xpos[self.torso_body_id]
gx = float(cur_torso[0] + stand_dist * np.cos(stand_heading))
gy = float(cur_torso[1] + stand_dist * np.sin(stand_heading))
gz = 1.168 # Grounded standing torso height equilibrium
self.getup_goal_pos = np.array([gx, gy, gz], dtype=np.float32)
self.goal_pos = self.getup_goal_pos.copy()
def _sample_new_goal(self, is_initial: bool = False):
"""Samples a new random navigation target within the 50m x 50m flat arena."""
torso_pos = self.data.xpos[self.torso_body_id]
if is_initial:
# Initial target: distance 5.0m - 7.5m in forward cone (-60 deg to +60 deg)
dist = float(np.random.uniform(5.0, 7.5))
angle = float(np.random.uniform(-np.pi / 3, np.pi / 3))
self.goal_transition_steps = 0 # Spawn directly facing target: immediate approach required
else:
# Subsequent target: distance 5.5m - 9.0m in full 360 deg omnidirectional bearing
dist = float(np.random.uniform(5.5, 9.0))
angle = float(np.random.uniform(-np.pi, np.pi))
self.goal_transition_steps = self.transition_leeway_window # ~1.85s direction change leeway
gx = float(np.clip(torso_pos[0] + dist * np.cos(angle), -22.0, 22.0))
gy = float(np.clip(torso_pos[1] + dist * np.sin(angle), -22.0, 22.0))
# Ensure newly sampled target is at least (arrival_radius + 1.5m = 3.5m) away from current torso
for _ in range(20):
if np.linalg.norm(np.array([gx - torso_pos[0], gy - torso_pos[1]])) >= (self.goal_arrival_radius + 1.5):
break
center_angle = float(np.arctan2(-torso_pos[1], -torso_pos[0]))
angle = float(center_angle + np.random.uniform(-np.pi / 4, np.pi / 4))
dist = float(np.random.uniform(5.5, 8.5))
gx = float(np.clip(torso_pos[0] + dist * np.cos(angle), -22.0, 22.0))
gy = float(np.clip(torso_pos[1] + dist * np.sin(angle), -22.0, 22.0))
gz = 1.25 # Torso/chest level
self.goal_pos = np.array([gx, gy, gz], dtype=np.float32)
self.min_dist_to_goal = float(np.linalg.norm(self.goal_pos - torso_pos))
self.prev_dist_to_goal = float(self.min_dist_to_goal)
self.steps_since_progress = 0
self.steps_not_approaching = 0
self.negative_vx_steps = 0
@property
def observation_dim(self) -> int:
"""
Total observation vector length:
- Torso height z: 1
- Torso orientation (forward & up vectors): 6
- Torso linear velocity (local/global): 3
- Torso angular velocity: 3
- 21 joint angles: 21
- 21 joint velocities: 21
- 2 foot contact forces: 2
- 512 solid-state LiDAR range rays (16x32 grid): 512
- 3 ego-centric relative waypoint vector (Ξ”x, Ξ”y, Ξ”z local): 3
- 3 metabolic stats (stamina, health, hunger): 3
Total: 1 + 6 + 3 + 3 + 21 + 21 + 2 + 512 + 3 + 3 = 575
"""
return 575
@property
def action_dim(self) -> int:
return self.n_actuators # 21
def reset(self, seed: Optional[int] = None, spawn_x: float = 0.0, initial_vx: float = 0.0) -> np.ndarray:
if seed is not None:
np.random.seed(seed)
mujoco.mj_resetData(self.model, self.data)
# Initial root pose: at X=spawn_x, Y=0, Z=1.33 (athletic bent-knee stance)
self.data.qpos[0] = spawn_x # X
self.data.qpos[1] = 0.0 # Y
self.data.qpos[2] = 1.33 # Z
self.data.qpos[3:7] = np.array([1.0, 0.0, 0.0, 0.0]) # Quaternion (unit)
# Staggered lead leg initialization: alternates 50% left lead, 50% right lead
# This breaks single-leg asymmetric gait traps and forces both thighs to cycle
lead_leg = np.random.choice(["right", "left"])
if lead_leg == "right":
self.data.qpos[self.r_hip_qpos] = -0.35 # Right thigh poised forward
self.data.qpos[self.r_knee_qpos] = 0.45
self.data.qpos[self.r_ankle_qpos] = -0.15
self.data.qpos[self.l_hip_qpos] = 0.05 # Left leg supporting in stance
self.data.qpos[self.l_knee_qpos] = 0.25
self.data.qpos[self.l_ankle_qpos] = -0.05
else:
self.data.qpos[self.l_hip_qpos] = -0.35 # Left thigh poised forward
self.data.qpos[self.l_knee_qpos] = 0.45
self.data.qpos[self.l_ankle_qpos] = -0.15
self.data.qpos[self.r_hip_qpos] = 0.05 # Right leg supporting in stance
self.data.qpos[self.r_knee_qpos] = 0.25
self.data.qpos[self.r_ankle_qpos] = -0.05
# Add slight stochastic noise to joint angles for exploration diversity
joint_noise = np.random.uniform(-0.03, 0.03, size=self.n_actuators)
self.data.qpos[7:] += joint_noise
# Zero velocities with tiny noise (or inject forward running momentum if pioneer spawn)
self.data.qvel[:] = np.random.uniform(-0.02, 0.02, size=self.n_qvel)
if initial_vx > 0.0:
self.data.qvel[0] = initial_vx + np.random.uniform(-0.1, 0.1)
# Reset internal stats
self.stamina = 1.0
self.health = 1.0
self.hunger = 0.0
self.milestones_reached = set()
self.waypoints_reached = set()
self.current_wp_idx = 0
if self.stream_mode == "stand":
# Stream 1: Quiet Standing Balance (Canonical Reference Equilibrium Stance)
self.data.qpos[0] = 0.0
self.data.qpos[1] = 0.0
self.data.qpos[2] = 1.168
self.data.qpos[3:7] = np.array([1.0, 0.0, 0.0, 0.0])
self.data.qvel[:] = 0.0
self.data.qpos[7:7 + self.n_actuators] = STAND_REF_QPOS.copy()
# Add slight stochastic exploration noise (+/- 0.015 rad ~ 0.8 deg) so agent learns to balance
self.data.qpos[7:7 + self.n_actuators] += np.random.uniform(-0.015, 0.015, size=self.n_actuators)
# Settle foot ground contacts smoothly (10 passive sub-steps with zero actuation)
mujoco.mj_forward(self.model, self.data)
for _ in range(10):
self.data.ctrl[:] = 0.0
mujoco.mj_step(self.model, self.data)
self.data.qvel[:] = 0.0
mujoco.mj_forward(self.model, self.data)
self.goal_pos = np.array([0.0, 0.0, 1.168], dtype=np.float32)
self.stand_still_steps = 0
self.stand_drift_steps = 0
self.stand_one_leg_steps = 0
self.goals_completed = 0
elif self.stream_mode == "getup":
# Stream 2: Fall Recovery (Randomized 360Β° poses & standing destination goal)
self._apply_fall_state()
self.prev_z = float(self.data.xpos[self.torso_body_id][2])
self.prev_head_z = float(self.data.xpos[self.head_body_id][2])
self.consecutive_upright_steps = 0
self.getup_stagnant_steps = 0
self.max_elevation_achieved = self.prev_z
self.max_head_elevation = self.prev_head_z
self.getup_success_count = 0
self.goals_completed = 0
elif self.stream_mode == "reach":
# Stream 2: Bilateral Hand Reaching (Bolted Torso with neutral resting arm configuration)
self._lock_torso_for_reach()
# Initialize both arms down at sides (natural resting posture at hip)
self.data.qpos[self.arm_qpos_indices[0]] = 0.65 + float(np.random.uniform(-0.03, 0.03)) # right_shoulder_x (down at side)
self.data.qpos[self.arm_qpos_indices[1]] = float(np.random.uniform(-0.03, 0.03)) # right_shoulder_y
self.data.qpos[self.arm_qpos_indices[2]] = float(np.random.uniform(-0.03, 0.03)) # right_elbow
self.data.qpos[self.arm_qpos_indices[3]] = -0.65 + float(np.random.uniform(-0.03, 0.03)) # left_shoulder_x (down at side)
self.data.qpos[self.arm_qpos_indices[4]] = float(np.random.uniform(-0.03, 0.03)) # left_shoulder_y
self.data.qpos[self.arm_qpos_indices[5]] = float(np.random.uniform(-0.03, 0.03)) # left_elbow
self.data.qvel[self.arm_dof_indices] = 0.0
mujoco.mj_forward(self.model, self.data)
self._sample_reach_target()
is_right_target = (self.active_reach_hand == "right")
active_h = self.data.geom_xpos[self.r_hand_geom_id] if is_right_target else self.data.geom_xpos[self.l_hand_geom_id]
self.prev_reach_dist = float(np.linalg.norm(active_h - self.reach_target_pos))
self.min_reach_dist = self.prev_reach_dist
self.reach_touches = 0
self.steps_on_reach_target = 0
self.steps_receding_reach = 0
self.goals_completed = 0
elif self.arena_mode:
self.goals_completed = 0
self._sample_new_goal(is_initial=True)
else:
# Fast-forward waypoint index if spawning downfield
for idx, (wp_pos, _, _) in enumerate(COURSE_WAYPOINTS):
if wp_pos[0] > spawn_x + 0.3:
self.current_wp_idx = idx
break
self.spawn_x = float(spawn_x)
self.furthest_x = float(spawn_x)
self.current_step = 0
self.steps_since_progress = 0
self.last_progress_x = float(spawn_x)
self.negative_vx_steps = 0
self.steps_not_approaching = 0
self.r_foot_contact_steps = 0
self.l_foot_contact_steps = 0
mujoco.mj_forward(self.model, self.data)
return self._get_obs()
def _get_obs(self) -> np.ndarray:
# 1. Torso state
torso_xpos = self.data.xpos[self.torso_body_id] # (3,)
torso_xmat = self.data.xmat[self.torso_body_id].reshape(3, 3) # (3, 3)
forward_vec = torso_xmat[:, 0] # Local X axis in world frame
up_vec = torso_xmat[:, 2] # Local Z axis in world frame
torso_cvel = self.data.cvel[self.torso_body_id] # 6D: [ang_vel(3), lin_vel(3)]
ang_vel = torso_cvel[:3]
lin_vel = torso_cvel[3:]
# 2. Joint states (21 angles, 21 velocities)
joint_qpos = self.data.qpos[7:7 + self.n_actuators]
joint_qvel = self.data.qvel[6:6 + self.n_actuators]
# 3. Touch sensors & Task-Specific Contact Indicators
touch_forces = np.zeros(2, dtype=np.float32)
if self.stream_mode == "reach":
# For isolated Reach, legs are bolted stationary. Repurpose touch_forces as explicit Active Hand Indicator:
# [1.0, 0.0] -> Right Hand Active | [0.0, 1.0] -> Left Hand Active
if self.active_reach_hand == "right":
touch_forces[0] = 1.0
touch_forces[1] = 0.0
else:
touch_forces[0] = 0.0
touch_forces[1] = 1.0
else:
if self.model.nsensor >= 2:
touch_forces[0] = np.clip(self.data.sensordata[0] / 100.0, 0.0, 5.0) # Right foot
touch_forces[1] = np.clip(self.data.sensordata[1] / 100.0, 0.0, 5.0) # Left foot
# 4. True Solid-State Angular LiDAR (512 rays from visor/eyes)
lidar_obs = self._compute_lidar_vision()
# 5. Dynamic Goal / Progressive Waypoint (Ego-centric local relative vector)
if self.stream_mode == "stand":
rel_world = self.goal_pos - torso_xpos
rel_local = torso_xmat.T @ rel_world
waypoint_obs = np.array([
np.clip(rel_local[0] / 10.0, -1.0, 1.0),
np.clip(rel_local[1] / 5.0, -1.0, 1.0),
np.clip(rel_local[2] / 2.0, -1.0, 1.0),
], dtype=np.float32)
elif self.stream_mode == "getup":
rel_world = self.getup_goal_pos - torso_xpos
rel_local = torso_xmat.T @ rel_world
waypoint_obs = np.array([
np.clip(rel_local[0] / 5.0, -1.0, 1.0),
np.clip(rel_local[1] / 5.0, -1.0, 1.0),
np.clip(rel_local[2] / 2.0, -1.0, 1.0),
], dtype=np.float32)
elif self.stream_mode == "reach":
# Direct closed-loop motor servoing: vector from ACTIVE HAND directly to the reach target orb!
active_h = self.data.geom_xpos[self.r_hand_geom_id] if (self.active_reach_hand == "right") else self.data.geom_xpos[self.l_hand_geom_id]
rel_hand_world = self.reach_target_pos - active_h
rel_hand_local = torso_xmat.T @ rel_hand_world
waypoint_obs = np.array([
np.clip(rel_hand_local[0] / 0.6, -1.0, 1.0),
np.clip(rel_hand_local[1] / 0.6, -1.0, 1.0),
np.clip(rel_hand_local[2] / 0.6, -1.0, 1.0),
], dtype=np.float32)
elif self.arena_mode:
target_pos = self.goal_pos
rel_world = target_pos - torso_xpos
rel_local = torso_xmat.T @ rel_world
waypoint_obs = np.array([
np.clip(rel_local[0] / 10.0, -1.0, 1.0),
np.clip(rel_local[1] / 5.0, -1.0, 1.0),
np.clip(rel_local[2] / 2.0, -1.0, 1.0),
], dtype=np.float32)
else:
target_pos = COURSE_WAYPOINTS[self.current_wp_idx][0]
rel_world = target_pos - torso_xpos
rel_local = torso_xmat.T @ rel_world
waypoint_obs = np.array([
np.clip(rel_local[0] / 10.0, -1.0, 1.0),
np.clip(rel_local[1] / 5.0, -1.0, 1.0),
np.clip(rel_local[2] / 2.0, -1.0, 1.0),
], dtype=np.float32)
# 6. Metabolic state (stamina, health, hunger/task-channel)
if self.stream_mode == "reach":
hand_indicator = 1.0 if (self.active_reach_hand == "right") else -1.0
metabolic_obs = np.array([self.stamina, self.health, hand_indicator], dtype=np.float32)
else:
metabolic_obs = np.array([self.stamina, self.health, self.hunger], dtype=np.float32)
obs = np.concatenate([
np.array([torso_xpos[2]], dtype=np.float32), # Torso Z (1)
forward_vec.astype(np.float32), # 3
up_vec.astype(np.float32), # 3
lin_vel.astype(np.float32), # 3
ang_vel.astype(np.float32), # 3
joint_qpos.astype(np.float32), # 21
joint_qvel.astype(np.float32), # 21
touch_forces, # 2
lidar_obs, # 512
waypoint_obs, # 3
metabolic_obs, # 3
])
return np.nan_to_num(obs, nan=0.0, posinf=50.0, neginf=-50.0)
def _compute_lidar_vision(self) -> np.ndarray:
"""
True Solid-State Angular LiDAR scan from the robot's visor/eyes.
Casts 128 rays (8 vertical elevation x 16 horizontal azimuth) using MuJoCo's native C raycaster.
Returns normalized range vector (128,) in [0.0, 1.0].
"""
head_pos = self.data.xpos[self.head_body_id]
head_mat = self.data.xmat[self.head_body_id].reshape(3, 3)
# Ray origin: 0.15m forward from head center through visor
ray_origin = np.ascontiguousarray(head_pos + head_mat @ self._visor_offset, dtype=np.float64)
# Transform all 128 precomputed local vectors to world coordinates
# (128, 3) @ (3, 3).T -> (128, 3)
world_vectors = self.lidar_local_vectors @ head_mat.T
normalized_ranges = np.zeros(self.num_lidar_rays, dtype=np.float32)
raw_grid = np.zeros((self.num_lidar_v, self.num_lidar_h), dtype=np.float32)
ray_idx = 0
for row in range(self.num_lidar_v):
for col in range(self.num_lidar_h):
vec = np.ascontiguousarray(world_vectors[ray_idx], dtype=np.float64)
dist = mujoco.mj_ray(self.model, self.data, ray_origin, vec, None, 1, -1, self._geomid_buf)
if dist >= 0:
gid = int(self._geomid_buf[0])
if gid in self.humanoid_geoms:
actual_dist = self.lidar_max_range
else:
actual_dist = min(dist, self.lidar_max_range)
else:
actual_dist = self.lidar_max_range
raw_grid[row, col] = actual_dist
# Zero-centered range: -0.5 (impact / 0m) to +0.5 (clear / 10m), eliminates DC bias
normalized_ranges[ray_idx] = (actual_dist / self.lidar_max_range) - 0.5
ray_idx += 1
self.current_lidar_grid = raw_grid
return normalized_ranges
def _compute_pseudo_vision(self, agent_x: float, agent_y: float, agent_z: float) -> np.ndarray:
"""Backwards compatibility fallback for 1D pseudo-vision."""
return self._compute_lidar_vision()[:16]
def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, bool, Dict]:
self.current_step += 1
# Clip action to [-1.0, 1.0]
action = np.clip(action, -1.0, 1.0)
self.data.ctrl[:] = action
if self.stream_mode == "reach":
# For isolated arm reaching, freeze motor commands to core and legs
self.data.ctrl[0:15] = 0.0
# Substep physics
for _ in range(self.frame_skip):
mujoco.mj_step(self.model, self.data)
if self.stream_mode == "reach":
self._lock_torso_for_reach()
# Extract state measurements
torso_xpos = self.data.xpos[self.torso_body_id]
x, y, z = torso_xpos[0], torso_xpos[1], torso_xpos[2]
torso_cvel = self.data.cvel[self.torso_body_id]
vx = torso_cvel[3] # Forward linear velocity
vy = torso_cvel[4]
torso_xmat = self.data.xmat[self.torso_body_id].reshape(3, 3)
up_z = torso_xmat[2, 2] # Cosine of torso vertical tilt
# Update metabolic stats
torque_cost = np.sum(np.square(action))
self.stamina = np.clip(self.stamina - (torque_cost * 0.0003) + 0.0005, 0.0, 1.0)
self.hunger = np.clip(self.hunger + 0.0001, 0.0, 1.0)
# Multi-limb ground contact detection for floor righting and locomotion
contact_r = False
contact_l = False
contact_hand = False
contact_knee = False
for c_idx in range(self.data.ncon):
c = self.data.contact[c_idx]
g1, g2 = c.geom1, c.geom2
if (g1 == self.r_foot_geom_id and g2 in self.world_obstacle_geoms) or (g2 == self.r_foot_geom_id and g1 in self.world_obstacle_geoms):
contact_r = True
if (g1 == self.l_foot_geom_id and g2 in self.world_obstacle_geoms) or (g2 == self.l_foot_geom_id and g1 in self.world_obstacle_geoms):
contact_l = True
if (g1 in self.support_hand_geoms and g2 in self.world_obstacle_geoms) or (g2 in self.support_hand_geoms and g1 in self.world_obstacle_geoms):
contact_hand = True
if (g1 in self.support_knee_geoms and g2 in self.world_obstacle_geoms) or (g2 in self.support_knee_geoms and g1 in self.world_obstacle_geoms):
contact_knee = True
touch_r = contact_r or (bool(self.data.sensordata[0] > 5.0) if self.model.nsensor >= 1 else False)
touch_l = contact_l or (bool(self.data.sensordata[1] > 5.0) if self.model.nsensor >= 2 else False)
if touch_r:
self.r_foot_contact_steps += 1
if touch_l:
self.l_foot_contact_steps += 1
upright_scale = float(np.clip((z - 0.45) / 0.40, 0.0, 1.0)) # 0 at z<=0.45m, 1.0 at z>=0.85m
in_leeway = False
torso_speed = float(np.linalg.norm(torso_cvel[3:]))
# ── STREAM 1: QUIET STANDING STILLNESS (REFERENCE POSE SIMILARITY) ────────
if self.stream_mode == "stand":
head_z = float(self.data.xpos[self.head_body_id][2])
dist_from_origin = float(np.linalg.norm([x, y]))
# 1. Joint Configuration Similarity:
# Gaussian kernel comparing 21 actuated joints against canonical STAND_REF_QPOS
current_qpos = self.data.qpos[7:7 + self.n_actuators]
diff_q = current_qpos - STAND_REF_QPOS
weighted_q_mse = float(np.sum(STAND_JOINT_WEIGHTS * (diff_q ** 2)) / np.sum(STAND_JOINT_WEIGHTS))
r_pose = float(np.exp(-3.0 * weighted_q_mse))
# 2. Torso Center-of-Mass Height Similarity (centered at z* = 1.168m):
r_height = float(np.exp(-30.0 * ((z - 1.168) ** 2)))
# 3. Torso Upright Orientation Similarity:
r_upright = float(np.exp(-12.0 * (max(0.0, 1.0 - up_z) ** 2)))
# 4. Multi-Rigid-Body Stillness Similarity (root linear, root angular, joint vel):
lin_vel_sq = float(np.sum(torso_cvel[3:] ** 2))
ang_vel_sq = float(np.sum(torso_cvel[:3] ** 2))
joint_vel_sq = float(np.mean(self.data.qvel[6:6 + self.n_actuators] ** 2))
r_stillness = float(np.exp(-4.0 * (lin_vel_sq + 0.5 * ang_vel_sq + 0.1 * joint_vel_sq)))
# 5. Foot Ground Contact Multiplier:
both_feet_down = bool(touch_r and touch_l)
one_foot_down = bool(touch_r or touch_l)
if both_feet_down:
m_feet = 1.0
elif one_foot_down:
m_feet = 0.45
else:
m_feet = 0.05
# 6. Action Control / Effort Penalty:
r_ctrl = -0.002 * float(np.sum(action ** 2))
# Bounded Gaussian Similarity Reward (Peak = 14.0/step at quiet standing equilibrium)
total_reward = m_feet * (4.0 * r_pose + 3.0 * r_height + 3.0 * r_upright + 4.0 * r_stillness) + r_ctrl
# Track stillness success criteria:
is_stand_still = bool(
r_pose > 0.85 and
r_height > 0.88 and
r_upright > 0.90 and
r_stillness > 0.80 and
both_feet_down
)
if is_stand_still:
self.stand_still_steps += 1
else:
self.stand_still_steps = max(0, self.stand_still_steps - 1)
# Check collision with environmental obstacles/ground using body geoms (head, chest, torso, visor)
is_obstacle_or_ground_impact = False
for c_idx in range(self.data.ncon):
c = self.data.contact[c_idx]
if (c.geom1 in self.body_impact_geoms and c.geom2 in self.world_obstacle_geoms) or \
(c.geom2 in self.body_impact_geoms and c.geom1 in self.world_obstacle_geoms):
is_obstacle_or_ground_impact = True
break
terminated = False
# Clean early termination only when genuinely fallen or severely displaced:
if z < 0.85 or up_z < 0.55 or is_obstacle_or_ground_impact:
terminated = True
total_reward -= 15.0
elif dist_from_origin > 1.20:
terminated = True
total_reward -= 10.0
# Mastery time limit: 384 steps (~9.6s)
truncated = (self.current_step >= 384)
if truncated and (not terminated) and z > 1.10 and both_feet_down:
total_reward += 50.0 # Mastery bonus for surviving 384 steps upright on both feet
# Detailed metrics for logging & dashboard
torso_roll = abs(float(torso_xmat[2, 1]))
torso_pitch = abs(float(torso_xmat[2, 0]))
spine_lateral = abs(float(self.data.qpos[self.abdomen_x_qpos]))
spine_sagittal = abs(float(self.data.qpos[self.abdomen_y_qpos]))
info = {
"x_pos": float(x), "y_pos": float(y), "z_pos": float(z),
"vx": float(vx), "torso_speed": float(torso_speed),
"furthest_x": float(self.furthest_x), "milestones": 0,
"waypoint_idx": 0, "waypoint_name": "Quiet Stand Stillness",
"dist_to_wp": float(dist_from_origin), "waypoints_completed": int(self.stand_still_steps > 150),
"goals_completed": int(self.stand_still_steps > 150),
"dist_to_goal": float(dist_from_origin),
"goal_arrival_radius": 1.20,
"goal_pos": [0.0, 0.0, 1.168],
"target_local_pos": [0.0, 0.0, 0.0],
"gaze_alignment": 1.0,
"stamina": float(self.stamina), "health": float(self.health), "hunger": float(self.hunger),
"step": self.current_step, "lidar_grid": self.current_lidar_grid.tolist(),
"in_leeway": False, "leeway_remaining": 0,
"stream_mode": "stand", "active_reach_hand": "none", "reach_touches": 0,
"getup_success": False,
"both_feet_down": bool(both_feet_down),
"one_foot_down": bool(one_foot_down),
"stand_still_score": float(is_stand_still),
"torso_roll": float(torso_roll),
"torso_pitch": float(torso_pitch),
"spine_lateral": float(spine_lateral),
"spine_sagittal": float(spine_sagittal),
"r_pose": float(r_pose),
"r_height": float(r_height),
"r_upright": float(r_upright),
"r_stillness": float(r_stillness),
}
return self._get_obs(), float(total_reward), terminated, truncated, info
# ── STREAM 2: FALL RECOVERY (GET-UP) ────────────────────────────────
elif self.stream_mode == "getup":
head_z = float(self.data.xpos[self.head_body_id][2])
delta_head_z = head_z - self.prev_head_z
self.prev_head_z = head_z
r_foot_z = float(self.data.xpos[self.r_foot_id][2])
l_foot_z = float(self.data.xpos[self.l_foot_id][2])
max_foot_z = max(r_foot_z, l_foot_z)
# 1. REWARD FOR HEAD & TORSO ELEVATION (Rises vertically, head leads above torso):
reward_head_elevation = 16.0 * delta_head_z
reward_torso_elevation = 12.0 * max(0.0, z - self.prev_z)
reward_head_height = 5.0 * max(0.0, head_z - 0.40)
reward_head_above_torso = 3.5 * max(0.0, head_z - z)
reward_upright = 3.0 * max(0.0, up_z - 0.30)
# 2. MULTI-PHASE GROUND CONTACT SUPPORT (Hands, Knees, Feet):
both_feet_down = touch_r and touch_l
one_foot_down = touch_r or touch_l
# Floor bracing support (head_z < 1.15m): reward hands pushing floor & knees tucked
reward_hand_brace = 4.0 if (contact_hand and head_z < 1.15) else 0.0
reward_knee_brace = 3.0 if (contact_knee and head_z < 1.15) else 0.0
# Standing stage (head_z >= 1.05m): reward planting feet
if both_feet_down:
reward_feet_ground = 6.0
elif one_foot_down:
reward_feet_ground = 2.0
elif head_z >= 1.20:
reward_feet_ground = -1.0 # Only penalize missing foot contact once upright
else:
reward_feet_ground = 0.0 # Zero penalty while floor-bound rolling or pushing up!
# Floor proximity for feet during ground phase
reward_feet_down = 2.0 * max(0.0, 0.30 - max_foot_z) if head_z < 1.15 else 0.0
# 3. MITIGATE LEG AIR-FLAILING:
# Only penalize feet higher than head IF head is actually on the ground and feet are violently kicked up
feet_above_head = max(0.0, max_foot_z - max(0.40, head_z))
penalty_shooting_legs = -4.0 * feet_above_head
# Air-flail penalty only when feet are high in air with zero contact
penalty_air_flail = -2.0 * max(0.0, max_foot_z - 0.50) if (head_z < 1.0 and not (contact_hand or contact_knee)) else 0.0
# 4. Energy cost & Arm Windmilling Suppression
arm_flail_penalty = -0.001 * float(np.sum(np.square(action[15:21])))
reward_ctrl = -0.001 * torque_cost + arm_flail_penalty
dist_to_stand_goal = float(np.linalg.norm(self.getup_goal_pos[:2] - np.array([x, y])))
reward_goal_heading = 2.0 * max(0.0, 1.0 - dist_to_stand_goal / 2.0) if head_z > 1.30 else 0.0
total_reward = (
reward_head_elevation +
reward_torso_elevation +
reward_head_height +
reward_head_above_torso +
reward_upright +
reward_hand_brace +
reward_knee_brace +
reward_feet_ground +
reward_feet_down +
penalty_shooting_legs +
penalty_air_flail +
reward_ctrl +
reward_goal_heading
)
self.prev_z = z
# Track peak head elevation for stagnation monitoring
if head_z > self.max_head_elevation:
self.max_head_elevation = head_z
self.getup_stagnant_steps = 0
else:
self.getup_stagnant_steps += 1
# Upright standing with feet planted verification:
# Requires head up (> 1.35m), torso upright (z > 1.05m, up_z > 0.70), and foot contact
is_upright = (head_z > 1.35) and (z > 1.05) and (up_z > 0.70) and (one_foot_down or (head_z > 1.45 and z > 1.10))
if is_upright:
self.consecutive_upright_steps += 1
total_reward += 3.0
if both_feet_down:
total_reward += 3.0 # Extra jackpot bonus when BOTH feet are planted!
else:
self.consecutive_upright_steps = 0
# Lenient Ground Collapse Stagnation & Rise Stall Deaths:
# Tolerates 140 steps (~3.5s) struggling on ground without head elevation
# Tolerates 180 steps (~4.5s) coordinating knee-tuck and feet planting to reach upright stance
is_ground_stuck = (head_z < 0.50 and self.getup_stagnant_steps > 140)
is_stall_stuck = (head_z >= 0.50 and self.consecutive_upright_steps == 0 and self.getup_stagnant_steps > 180)
terminated = False
if self.consecutive_upright_steps >= 12:
total_reward += 100.0 # Upright Recovery Jackpot!
self.getup_success_count += 1
terminated = True
elif is_ground_stuck or is_stall_stuck:
terminated = True
total_reward -= 10.0 # Timeout stagnation death
truncated = (self.current_step >= 280)
target_pos = self.getup_goal_pos
info = {
"x_pos": float(x), "y_pos": float(y), "z_pos": float(z),
"vx": float(vx), "torso_speed": float(torso_speed),
"furthest_x": float(self.furthest_x), "milestones": 0,
"waypoint_idx": 0, "waypoint_name": f"Get-Up ({self.fall_state_type})",
"dist_to_wp": float(max(0.0, 1.55 - head_z)),
"waypoints_completed": int(self.consecutive_upright_steps >= 12),
"goals_completed": int(self.consecutive_upright_steps >= 12),
"dist_to_goal": float(dist_to_stand_goal),
"goal_arrival_radius": 0.50,
"goal_pos": self.getup_goal_pos.tolist(),
"target_local_pos": (torso_xmat.T @ (target_pos - torso_xpos)).tolist(),
"gaze_alignment": float(max(0.0, up_z)),
"stamina": float(self.stamina), "health": float(self.health), "hunger": float(self.hunger),
"step": self.current_step, "lidar_grid": self.current_lidar_grid.tolist(),
"in_leeway": False, "leeway_remaining": 0,
"stream_mode": "getup", "active_reach_hand": "none", "reach_touches": 0,
"getup_success": bool(self.consecutive_upright_steps >= 12),
"head_z": float(head_z),
"both_feet_down": bool(both_feet_down),
}
return self._get_obs(), float(total_reward), terminated, truncated, info
# ── STREAM 2: BILATERAL HAND REACHING (BOLTED BODY) ────────────────
elif self.stream_mode == "reach":
# Torso, core, and legs are bolted in neutral athletic upright stance.
# Control is isolated strictly to the 6 arm actuators reaching for 3D orbs.
r_h = self.data.geom_xpos[self.r_hand_geom_id]
l_h = self.data.geom_xpos[self.l_hand_geom_id]
dist_r = float(np.linalg.norm(r_h - self.reach_target_pos))
dist_l = float(np.linalg.norm(l_h - self.reach_target_pos))
# The designated active hand is the hand closer to the target ball
is_right_target = (self.active_reach_hand == "right")
if is_right_target:
active_h = r_h
wrong_h = l_h
dist_to_target = dist_r
dist_inactive = dist_l
active_body_id = self.r_arm_body_id
wrong_body_id = self.l_arm_body_id
active_act = action[15:18]
wrong_act = action[18:21]
active_qvel = self.data.qvel[self.arm_dof_indices[0:3]]
wrong_qvel = self.data.qvel[self.arm_dof_indices[3:6]]
wrong_qpos = self.data.qpos[self.arm_qpos_indices[3:6]]
wrong_rest_qpos = np.array([-0.65, 0.0, 0.0], dtype=np.float32)
else:
active_h = l_h
wrong_h = r_h
dist_to_target = dist_l
dist_inactive = dist_r
active_body_id = self.l_arm_body_id
wrong_body_id = self.r_arm_body_id
active_act = action[18:21]
wrong_act = action[15:18]
active_qvel = self.data.qvel[self.arm_dof_indices[3:6]]
wrong_qvel = self.data.qvel[self.arm_dof_indices[0:3]]
wrong_qpos = self.data.qpos[self.arm_qpos_indices[0:3]]
wrong_rest_qpos = np.array([0.65, 0.0, 0.0], dtype=np.float32)
# 1. Main reach distance progress and proximity reward
dist_prog = self.prev_reach_dist - dist_to_target
reward_reach_prog = 25.0 * dist_prog
reward_proximity = 6.0 * np.exp(-4.0 * dist_to_target)
# 2. MASSIVE SPEED EFFICIENCY REWARDS:
# a) Direct closing velocity vector along the unit ray connecting active hand to target orb
delta_active_target = self.reach_target_pos - active_h
dist_hand_target = max(1e-4, float(np.linalg.norm(delta_active_target)))
u_reach_target = delta_active_target / dist_hand_target
active_hand_vel = self.data.cvel[active_body_id][3:6]
v_closing = float(np.dot(active_hand_vel, u_reach_target))
# Speed efficiency bonus for thrusting active hand toward target
reward_closing_speed = 4.0 * max(0.0, v_closing)
# b) Mild latency penalty (encourages prompt reaching)
penalty_latency = -0.02 * (float(self.steps_on_reach_target) / 40.0)
# 3. Slight reward for moving the correct hand
active_hand_speed = float(np.linalg.norm(active_hand_vel))
active_qvel_mag = float(np.sum(np.abs(active_qvel)))
reward_correct_motion = 0.15 * min(1.0, active_hand_speed / 0.25) + 0.10 * np.tanh(active_qvel_mag)
# 4. REGULARIZATION ON INACTIVE HAND (stay neutral and relaxed at side)
wrong_hand_speed = float(np.linalg.norm(self.data.cvel[wrong_body_id][3:6]))
penalty_wrong_act = -0.05 * float(np.sum(np.square(wrong_act)))
penalty_wrong_qvel = -0.05 * float(np.sum(np.abs(wrong_qvel)))
penalty_wrong_speed = -0.20 * min(1.0, wrong_hand_speed)
penalty_wrong_raised = -0.50 * max(0.0, float(wrong_h[2]) - 1.28)
penalty_wrong_posture = -0.20 * min(1.0, float(np.linalg.norm(wrong_qpos - wrong_rest_qpos)))
penalty_wrong_lead = -0.50 * max(0.0, dist_to_target - dist_inactive)
penalty_wrong_encroach = -0.50 * max(0.0, 0.35 - dist_inactive)
penalty_wrong_touch = -2.0 if dist_inactive < 0.14 else 0.0
# General control penalty on active arm to encourage efficient articulation
reward_ctrl = -0.001 * float(np.sum(np.square(active_act)))
self.steps_on_reach_target += 1
if dist_to_target < self.min_reach_dist:
self.min_reach_dist = dist_to_target
self.steps_receding_reach = 0
elif dist_to_target > self.min_reach_dist + 0.25:
self.steps_receding_reach += 1
reward_touch = 0.0
speed_efficiency_jackpot = 0.0
if dist_to_target < 0.14:
# 1. Base touch award
base_touch = 40.0
# 2. Speed efficiency jackpot: maximum when acquired in few steps
speed_efficiency_jackpot = max(0.0, 110.0 - float(self.steps_on_reach_target)) * 0.40
# 3. Consecutive hit streak multiplier
touch_streak_bonus = min(30.0, float(self.reach_touches) * 6.0)
reward_touch = base_touch + speed_efficiency_jackpot + touch_streak_bonus
self.reach_touches += 1
self._sample_reach_target()
is_right_target = (self.active_reach_hand == "right")
active_h = self.data.geom_xpos[self.r_hand_geom_id] if is_right_target else self.data.geom_xpos[self.l_hand_geom_id]
dist_to_target = float(np.linalg.norm(active_h - self.reach_target_pos))
# Reset timers for next target!
self.steps_on_reach_target = 0
self.steps_receding_reach = 0
self.min_reach_dist = dist_to_target
self.prev_reach_dist = dist_to_target
total_reward = (
reward_reach_prog +
reward_proximity +
reward_closing_speed +
reward_correct_motion +
reward_touch +
penalty_latency +
penalty_wrong_act +
penalty_wrong_qvel +
penalty_wrong_speed +
penalty_wrong_raised +
penalty_wrong_posture +
penalty_wrong_lead +
penalty_wrong_encroach +
penalty_wrong_touch +
reward_ctrl
)
# Lenient Reaching Time Limits:
# Tolerates 110 steps (~2.75s) per target orb, or 45 steps receding away
is_reach_timeout = (self.steps_on_reach_target > 110)
is_reach_receding = (self.steps_receding_reach > 45)
terminated = False
if is_reach_timeout or is_reach_receding:
terminated = True
total_reward -= 2.0 # Mild timeout penalty
truncated = (self.current_step >= 384)
target_pos = self.reach_target_pos
head_xmat = self.data.xmat[self.head_body_id].reshape(3, 3)
head_fwd = head_xmat[:, 0]
delta_head_target = target_pos - self.data.xpos[self.head_body_id]
cos_gaze = float(np.dot(head_fwd, delta_head_target / max(1e-3, np.linalg.norm(delta_head_target))))
info = {
"x_pos": float(x), "y_pos": float(y), "z_pos": float(z),
"vx": float(vx), "torso_speed": float(torso_speed),
"furthest_x": float(self.furthest_x), "milestones": 0,
"waypoint_idx": 0, "waypoint_name": f"Reach ({self.active_reach_hand.upper()})",
"dist_to_wp": float(dist_to_target), "waypoints_completed": int(self.reach_touches),
"goals_completed": int(self.reach_touches),
"dist_to_goal": float(dist_to_target),
"goal_arrival_radius": 0.14,
"goal_pos": self.reach_target_pos.tolist(),
"target_local_pos": (torso_xmat.T @ (target_pos - torso_xpos)).tolist(),
"gaze_alignment": float(cos_gaze),
"stamina": float(self.stamina), "health": float(self.health), "hunger": float(self.hunger),
"step": self.current_step, "lidar_grid": self.current_lidar_grid.tolist(),
"in_leeway": False, "leeway_remaining": 0,
"stream_mode": "reach", "active_reach_hand": str(self.active_reach_hand),
"reach_touches": int(self.reach_touches), "getup_success": False, "stand_still_score": 0.0,
"wrong_hand_dist": float(dist_inactive), "wrong_hand_touch": bool(dist_inactive < 0.14),
"wrong_hand_speed": float(wrong_hand_speed),
"wrong_arm_action_penalty": float(penalty_wrong_act),
"wrong_arm_raised_penalty": float(penalty_wrong_raised),
"v_closing": float(v_closing),
"reward_closing_speed": float(reward_closing_speed),
"speed_efficiency_jackpot": float(speed_efficiency_jackpot),
}
return self._get_obs(), float(total_reward), terminated, truncated, info
# Target Navigation Vector & Forward Propulsion
if self.arena_mode:
target_pos = self.goal_pos
delta_target = target_pos - torso_xpos
dist_to_target = float(np.linalg.norm(delta_target))
dist_xy = float(np.linalg.norm(delta_target[:2]))
# Check dynamic goal arrival (entering wide arrival radius)
reward_arrival = 0.0
if dist_xy < self.goal_arrival_radius:
self.goals_completed += 1
# Escalating jackpot for chained consecutive goals: 50 for 1st, 85 for 2nd, 120 for 3rd, 155 for 4th...
reward_arrival = 50.0 + 35.0 * (self.goals_completed - 1)
self._sample_new_goal(is_initial=False)
target_pos = self.goal_pos
delta_target = target_pos - torso_xpos
dist_to_target = float(np.linalg.norm(delta_target))
dist_xy = float(np.linalg.norm(delta_target[:2]))
self.prev_dist_to_goal = dist_to_target
dist_delta = 0.0
else:
dist_delta = self.prev_dist_to_goal - dist_to_target
self.prev_dist_to_goal = dist_to_target
# Leeway period active after goal completion to allow decelerating & turning
in_leeway = (self.goal_transition_steps > 0)
if in_leeway:
self.goal_transition_steps -= 1
# While turning during leeway, continuously update closest distance baseline to current position
self.min_dist_to_goal = dist_to_target
# 1. Distance progress potential (scaled with goal streak) & sharp penalty for moving away
streak_mult = 1.0 + 0.35 * min(4, self.goals_completed)
if dist_delta >= 0.0:
reward_dist_progress = 3.0 * streak_mult * dist_delta
else:
# Double penalty for moving further away from target outside of turn leeway
reward_dist_progress = 0.0 if in_leeway else (6.0 * dist_delta)
# 2. Penalty for drifting/retreating beyond the closest distance achieved so far (muted in leeway)
retreat_distance = max(0.0, dist_to_target - self.min_dist_to_goal)
penalty_retreat = 0.0 if in_leeway else (-3.0 * retreat_distance if retreat_distance > 0.10 else 0.0)
unit_target_dir = delta_target / max(0.1, dist_to_target)
vel_towards_target = float(np.dot(torso_cvel[3:], unit_target_dir))
torso_speed = float(np.linalg.norm(torso_cvel[3:]))
# 3. Controlled Forward Movement vs Severe Retrograde Penalty
# Cap forward reward at 1.8 m/s so Olympic sprinting gives no extra reward
v_eff = min(max(0.0, vel_towards_target), 1.8)
reward_forward = 3.0 * v_eff * upright_scale
penalty_retrograde = 0.0 if (vel_towards_target >= 0.0 or in_leeway) else (5.0 * vel_towards_target)
# 4. Heading guidance velocity (steep negative slope when heading backwards outside leeway)
if vel_towards_target >= 0.0:
reward_heading = 1.5 * v_eff * upright_scale
else:
reward_heading = 0.0 if in_leeway else (3.0 * vel_towards_target)
# 5. Speed Regulation Penalty: Penalize reckless overspeeding (> 2.0 m/s) to promote stable locomotion
penalty_overspeed = 0.0
if torso_speed > 2.0:
penalty_overspeed = -3.0 * ((torso_speed - 2.0) ** 2)
# 6. Turn alignment & braking reward during transition leeway: encourage slowing down & rapid yaw alignment to new goal
reward_turn = 0.0
reward_brake = 0.0
if in_leeway:
torso_fwd = torso_xmat[:, 0]
cos_torso_goal = float(np.dot(torso_fwd[:2], unit_target_dir[:2]))
# Reward rotating body to face new target (+0.0 to +1.0)
reward_turn = 1.0 * max(0.0, cos_torso_goal)
# Reward decelerating below 1.6 m/s so the agent can pivot cleanly
if torso_speed < 1.6:
reward_brake = 0.6 * (1.6 - torso_speed)
# 7. Subsequent Goal Urgency: Time-decay pressure once Goal 1 is cleared
penalty_urgency = 0.0
if self.goals_completed >= 1 and not in_leeway:
penalty_urgency = -0.08 * min(3, self.goals_completed)
reward_wp = (
reward_arrival +
reward_dist_progress +
penalty_retreat +
penalty_retrograde +
penalty_overspeed +
reward_turn +
reward_brake +
penalty_urgency
)
reward_drift = 0.0
reward_milestone = 0.0
dist_to_wp = dist_to_target
else:
# Obstacle course mode
wp_pos, wp_name, wp_radius = COURSE_WAYPOINTS[self.current_wp_idx]
target_pos = wp_pos
delta_to_wp = wp_pos - torso_xpos
dist_to_wp = float(np.linalg.norm(delta_to_wp))
dist_xy = float(np.linalg.norm(delta_to_wp[:2]))
reward_forward = 3.5 * vx * upright_scale
if vx < -0.05:
reward_forward += 1.5 * vx
wp_reached = False
if dist_xy < wp_radius:
wp_reached = True
elif torso_xpos[0] >= wp_pos[0] and abs(torso_xpos[1] - wp_pos[1]) < (wp_radius + 0.3):
wp_reached = True
reward_wp = 0.0
if wp_reached and self.current_wp_idx not in self.waypoints_reached:
self.waypoints_reached.add(self.current_wp_idx)
if self.current_wp_idx == len(COURSE_WAYPOINTS) - 1:
reward_wp = 100.0
elif 3 <= self.current_wp_idx <= 14:
reward_wp = 25.0
else:
reward_wp = 15.0
if self.current_wp_idx < len(COURSE_WAYPOINTS) - 1:
self.current_wp_idx += 1
wp_pos, wp_name, wp_radius = COURSE_WAYPOINTS[self.current_wp_idx]
target_pos = wp_pos
delta_to_wp = wp_pos - torso_xpos
dist_to_wp = float(np.linalg.norm(delta_to_wp))
wp_unit_dir = delta_to_wp / max(0.1, dist_to_wp)
vel_towards_wp = float(np.dot(torso_cvel[3:], wp_unit_dir))
reward_heading = 2.0 * max(-0.5, min(vel_towards_wp, 3.5)) * upright_scale
target_y = float(wp_pos[1])
reward_drift = -0.5 * abs(y - target_y)
reward_milestone = 0.0
for threshold, name in STAGE_MILESTONES:
if x >= threshold and threshold not in self.milestones_reached:
self.milestones_reached.add(threshold)
reward_milestone += 50.0
# Head Gaze Alignment Reward: mild reward for facing objective with head/visor
head_xmat = self.data.xmat[self.head_body_id].reshape(3, 3)
head_fwd = head_xmat[:, 0] # Local X axis of head (forward gaze direction)
head_pos = self.data.xpos[self.head_body_id]
delta_head_target = target_pos - head_pos
dist_head_target = float(np.linalg.norm(delta_head_target))
u_head_target = delta_head_target / max(1e-3, dist_head_target)
cos_gaze = float(np.dot(head_fwd, u_head_target))
reward_gaze = 0.15 * max(0.0, cos_gaze)
# 2. Upright posture reward (staying balanced above 0.70m)
reward_upright = 1.0 if z > 0.70 else 0.0
# 3. Upright orientation bonus & active lateral roll stabilization
roll_y = torso_xmat[1, 2]
reward_roll = -2.0 * abs(roll_y)
reward_orient = 0.8 * max(0.0, up_z) + reward_roll
# Advance speed for biomechanical stepping gating
forward_vec = torso_xmat[:, 0]
v_advance = float(np.dot(torso_cvel[3:6], forward_vec)) if self.arena_mode else vx
# 4. Alternating Bipedal Stepping Cadence
reward_gait = 0.0
if z > 0.55 and v_advance > 0.10:
if touch_r ^ touch_l:
reward_gait = 2.0 * min(v_advance, 1.5)
elif touch_r and touch_l:
reward_gait = -0.4 * min(v_advance, 1.5)
# 5. Knee Articulation & Stance Stability
r_knee_q = self.data.qpos[self.r_knee_qpos]
l_knee_q = self.data.qpos[self.l_knee_qpos]
reward_knee = 0.0
if z > 0.55 and v_advance > 0.10:
if not touch_r:
if 0.25 <= r_knee_q <= 1.1:
reward_knee += 0.3
elif r_knee_q > 1.6:
reward_knee -= 0.5 * (r_knee_q - 1.6)
if not touch_l:
if 0.25 <= l_knee_q <= 1.1:
reward_knee += 0.3
elif l_knee_q > 1.6:
reward_knee -= 0.5 * (l_knee_q - 1.6)
if touch_r and r_knee_q < 0.50:
reward_knee += 0.2
if touch_l and l_knee_q < 0.50:
reward_knee += 0.2
# 6. Swing foot forward velocity bonus
r_foot_vfwd = float(np.dot(self.data.cvel[self.r_foot_id][3:6], forward_vec)) if self.arena_mode else self.data.cvel[self.r_foot_id][3]
l_foot_vfwd = float(np.dot(self.data.cvel[self.l_foot_id][3:6], forward_vec)) if self.arena_mode else self.data.cvel[self.l_foot_id][3]
reward_swing = 0.0
if z > 0.55 and v_advance > 0.10:
if not touch_r and r_foot_vfwd > 0.15:
reward_swing += 0.3 * min(r_foot_vfwd, 2.0)
if not touch_l and l_foot_vfwd > 0.15:
reward_swing += 0.3 * min(l_foot_vfwd, 2.0)
# 7. Thigh Forward Stride Separation
r_hip_q = self.data.qpos[self.r_hip_qpos]
l_hip_q = self.data.qpos[self.l_hip_qpos]
reward_stride = 0.0
if z > 0.55 and v_advance > 0.10:
if not touch_r:
stride_diff = l_hip_q - r_hip_q
if 0.1 < stride_diff < 1.3:
reward_stride = 1.0 * stride_diff
elif not touch_l:
stride_diff = r_hip_q - l_hip_q
if 0.1 < stride_diff < 1.3:
reward_stride = 1.0 * stride_diff
# 8. Bilateral Stance Symmetry (Duty-Cycle Equalizer)
reward_symmetry = 0.0
total_contacts = self.r_foot_contact_steps + self.l_foot_contact_steps
if self.current_step > 10 and total_contacts > 6 and z > 0.55:
asym = abs(self.r_foot_contact_steps - self.l_foot_contact_steps) / total_contacts
if asym > 0.20:
reward_symmetry = -2.0 * (asym - 0.20)
# 9. Arm Windmilling Suppression (Penalize flailing arms to force balance from core and legs)
arm_flail_penalty = -0.008 * float(np.sum(np.square(action[15:21])))
# 10. Energy expenditure penalty
reward_ctrl = -0.0015 * torque_cost + arm_flail_penalty
total_reward = (
reward_forward +
reward_heading +
reward_wp +
reward_gaze +
reward_upright +
reward_orient +
reward_gait +
reward_knee +
reward_swing +
reward_stride +
reward_symmetry +
reward_ctrl +
reward_drift +
reward_milestone
)
if self.arena_mode:
rad_dist = float(np.linalg.norm([x, y]))
if rad_dist > self.furthest_x:
self.furthest_x = rad_dist
if in_leeway:
# During turning leeway window: reset & do not accumulate moving-away or stagnation counters
self.steps_not_approaching = 0
self.negative_vx_steps = 0
self.steps_since_progress = 0
self.min_dist_to_goal = dist_to_target
else:
# Active approach check: requires advancing toward target (> 0.15 m/s and closing distance)
is_actively_approaching = (vel_towards_target >= 0.15) and (dist_delta > 0.001)
if is_actively_approaching:
self.steps_not_approaching = 0
else:
self.steps_not_approaching += 1
# Receding tracker: moving away from target
is_receding = (vel_towards_target < -0.05) or (dist_delta < -0.003)
if is_receding:
self.negative_vx_steps += 1
else:
self.negative_vx_steps = max(0, self.negative_vx_steps - 1)
# Record best approach distance
if dist_to_target < self.min_dist_to_goal:
self.min_dist_to_goal = dist_to_target
self.steps_since_progress = 0
else:
self.steps_since_progress += 1
else:
if x > self.furthest_x:
self.furthest_x = x
if x > self.last_progress_x + 0.06:
self.last_progress_x = x
self.steps_since_progress = 0
else:
self.steps_since_progress += 1
if vx < -0.15:
self.negative_vx_steps += 1
else:
self.negative_vx_steps = 0
# ── TERMINATION & TRUNCATION ──────────────────────────────────────────
# Exact preserved death conditions:
# 1. Physical collision: head, visor, chest, or upper torso directly impacts the floor/obstacles
# 2. True collapse: torso center of mass drops below 0.35m or head drops below 0.45m
# 3. Complete inversion: torso tilted > 81 deg from vertical (up_z < 0.15)
# 4. Arena boundary departure (|x| > 24m or |y| > 24m) or course lateral departure (|y| > 3.0m)
# 5. Fast elimination when moving away or not actively approaching target (outside leeway)
terminated = False
is_obstacle_or_ground_impact = False
for c_idx in range(self.data.ncon):
c = self.data.contact[c_idx]
g1, g2 = c.geom1, c.geom2
if (g1 in self.body_impact_geoms and g2 in self.world_obstacle_geoms) or \
(g2 in self.body_impact_geoms and g1 in self.world_obstacle_geoms):
is_obstacle_or_ground_impact = True
break
head_z = self.data.xpos[self.head_body_id][2]
is_altitude_collapse = (z < 0.35) or (head_z < 0.45)
is_tilt_collapse = (up_z < 0.15)
is_physical_fall = is_obstacle_or_ground_impact or is_altitude_collapse or is_tilt_collapse
if self.arena_mode:
is_boundary_departure = (abs(x) > 24.0 or abs(y) > 24.0)
# FAST KILL TRIGGERS (active outside turning leeway period):
# 1. Backing away kill: moved away from target for > 25 steps (~0.6s)
is_moving_away_kill = (not in_leeway) and (self.negative_vx_steps > 25)
# 2. Retreat distance kill: backed up > 0.40m further than closest approach achieved
is_retreat_kill = (not in_leeway) and (dist_to_target > self.min_dist_to_goal + 0.40) and (self.current_step > 15)
# 3. Not actively approaching kill: dawdling, circling, or stationary for > 80 steps (~2.0s)
is_not_approaching_kill = (not in_leeway) and (self.current_step > 25 and self.steps_not_approaching > 80)
# 4. Stagnation kill: no new record distance achieved for > 85 steps (~2.125s)
is_stagnant_kill = (not in_leeway) and (self.current_step > 25 and self.steps_since_progress > 85)
if is_physical_fall:
terminated = True
total_reward -= 20.0
self.health = max(0.0, self.health - 0.2)
elif is_boundary_departure:
terminated = True
total_reward -= 20.0
elif is_moving_away_kill or is_retreat_kill:
terminated = True
total_reward -= 25.0 # Severe penalty for retreating/moving away from target
elif is_not_approaching_kill or is_stagnant_kill:
terminated = True
total_reward -= 25.0 # Severe penalty for failing to actively approach target
# Dropped Goal Streak Penalty:
# If the agent achieved at least 1 goal, but died before reaching the next goal,
# apply a moderate penalty to encourage multi-goal chaining without negating arrival returns
if terminated and self.goals_completed >= 1:
dropped_goal_penalty = -15.0 * (1.0 + 0.2 * min(3, self.goals_completed - 1))
total_reward += dropped_goal_penalty
else:
is_moving_backwards = (x < self.spawn_x - 0.50) or (self.negative_vx_steps > 60)
is_stagnant = (self.current_step > 50 and self.steps_since_progress > 90)
is_lateral_departure = (abs(y) > 3.0)
if is_physical_fall:
terminated = True
total_reward -= 20.0
self.health = max(0.0, self.health - 0.2)
elif is_moving_backwards:
terminated = True
total_reward -= 15.0
elif is_stagnant:
terminated = True
total_reward -= 10.0
elif is_lateral_departure:
terminated = True
total_reward -= 20.0
if x >= 47.0:
terminated = True
total_reward += 200.0
truncated = (self.current_step >= self.max_steps)
info = {
"x_pos": float(x),
"y_pos": float(y),
"z_pos": float(z),
"vx": float(vx),
"torso_speed": float(torso_speed),
"furthest_x": float(self.furthest_x),
"milestones": len(self.milestones_reached),
"waypoint_idx": int(self.current_wp_idx),
"waypoint_name": str(COURSE_WAYPOINTS[self.current_wp_idx][1]) if not self.arena_mode else f"Arena Target #{self.goals_completed + 1}",
"dist_to_wp": float(dist_to_wp),
"waypoints_completed": len(self.waypoints_reached) if not self.arena_mode else int(self.goals_completed),
"goals_completed": int(self.goals_completed),
"dist_to_goal": float(dist_to_wp),
"goal_arrival_radius": float(self.goal_arrival_radius) if self.arena_mode else 1.0,
"goal_pos": self.goal_pos.tolist() if self.arena_mode else list(COURSE_WAYPOINTS[self.current_wp_idx][0]),
"target_local_pos": (torso_xmat.T @ (target_pos - torso_xpos)).tolist(),
"gaze_alignment": float(cos_gaze),
"stamina": float(self.stamina),
"health": float(self.health),
"hunger": float(self.hunger),
"step": self.current_step,
"lidar_grid": self.current_lidar_grid.tolist(),
"in_leeway": bool(in_leeway) if self.arena_mode else False,
"leeway_remaining": int(self.goal_transition_steps) if self.arena_mode else 0,
"stream_mode": "cruise",
"active_reach_hand": "none",
"reach_touches": 0,
"getup_success": False,
"stand_still_score": 0.0,
}
return self._get_obs(), float(total_reward), terminated, truncated, info
def render_frame(
self,
width: int = 480,
height: int = 320,
camera_name: str = "behind_cam",
ghost_datas: Optional[List[Any]] = None,
ghost_alpha: Optional[float] = None,
) -> np.ndarray:
"""
Renders an RGB numpy image of the current simulation frame.
Optionally injects translucent ghost avatars into mjvScene for parallel multi-agent swarm visualization.
"""
if self._renderer is None or self._renderer.width != width or self._renderer.height != height:
self._renderer = mujoco.Renderer(self.model, height=height, width=width)
self._renderer.update_scene(self.data, camera=camera_name)
# Inject dynamic course waypoint goal markers into mjvScene
scene = self._renderer.scene
identity_mat = np.eye(3).flatten()
if self.stream_mode == "reach":
# Stream 3 Reach: Glowing magenta hand target orb + vertical guide line
rx, ry, rz = float(self.reach_target_pos[0]), float(self.reach_target_pos[1]), float(self.reach_target_pos[2])
if scene.ngeom < scene.maxgeom:
dst_orb = scene.geoms[scene.ngeom]
orb_pos = np.array([rx, ry, rz], dtype=np.float64)
orb_size = np.array([0.09, 0.09, 0.09], dtype=np.float64)
orb_rgba = np.array([0.88, 0.20, 0.95, 0.95], dtype=np.float32) # Radiant Magenta
mujoco.mjv_initGeom(dst_orb, mujoco.mjtGeom.mjGEOM_SPHERE, orb_size, orb_pos, identity_mat, orb_rgba)
dst_orb.transparent = 1
scene.ngeom += 1
if scene.ngeom < scene.maxgeom:
dst_guide = scene.geoms[scene.ngeom]
guide_pos = np.array([rx, ry, max(0.01, rz / 2.0)], dtype=np.float64)
guide_size = np.array([0.012, max(0.01, rz / 2.0), 0.0], dtype=np.float64)
guide_rgba = np.array([0.88, 0.20, 0.95, 0.35], dtype=np.float32)
mujoco.mjv_initGeom(dst_guide, mujoco.mjtGeom.mjGEOM_CYLINDER, guide_size, guide_pos, identity_mat, guide_rgba)
dst_guide.transparent = 1
scene.ngeom += 1
elif self.stream_mode == "getup":
# Stream 2 Get-Up: Golden upright elevation target disc
torso_xpos = self.data.xpos[self.torso_body_id]
tx, ty = float(torso_xpos[0]), float(torso_xpos[1])
if scene.ngeom < scene.maxgeom:
dst_target = scene.geoms[scene.ngeom]
target_pos = np.array([tx, ty, 1.33], dtype=np.float64)
target_size = np.array([0.25, 0.02, 0.0], dtype=np.float64)
target_rgba = np.array([0.98, 0.75, 0.15, 0.65], dtype=np.float32) # Glowing Amber
mujoco.mjv_initGeom(dst_target, mujoco.mjtGeom.mjGEOM_CYLINDER, target_size, target_pos, identity_mat, target_rgba)
dst_target.transparent = 1
scene.ngeom += 1
elif self.stream_mode == "stand":
# Stream 1 Stand: Cyan stillness equilibrium ground pad
if scene.ngeom < scene.maxgeom:
dst_ring = scene.geoms[scene.ngeom]
ring_pos = np.array([0.0, 0.0, 0.015], dtype=np.float64)
ring_size = np.array([0.35, 0.015, 0.0], dtype=np.float64)
ring_rgba = np.array([0.22, 0.74, 0.97, 0.50], dtype=np.float32) # Cyan Pad
mujoco.mjv_initGeom(dst_ring, mujoco.mjtGeom.mjGEOM_CYLINDER, ring_size, ring_pos, identity_mat, ring_rgba)
dst_ring.transparent = 1
scene.ngeom += 1
elif self.arena_mode:
# ACTIVE ARENA GOAL: Glowing floating target orb + vertical laser beacon column + ground target ring
gx, gy, gz = float(self.goal_pos[0]), float(self.goal_pos[1]), float(self.goal_pos[2])
# 1. Floating target orb (sphere)
if scene.ngeom < scene.maxgeom:
dst_orb = scene.geoms[scene.ngeom]
orb_pos = np.array([gx, gy, gz + 0.35], dtype=np.float64)
orb_size = np.array([0.22, 0.22, 0.22], dtype=np.float64)
orb_rgba = np.array([0.98, 0.75, 0.15, 0.90], dtype=np.float32) # Vibrant Gold
mujoco.mjv_initGeom(dst_orb, mujoco.mjtGeom.mjGEOM_SPHERE, orb_size, orb_pos, identity_mat, orb_rgba)
dst_orb.transparent = 1
scene.ngeom += 1
# 2. Vertical laser beacon column (cylinder)
if scene.ngeom < scene.maxgeom:
dst_beam = scene.geoms[scene.ngeom]
beam_pos = np.array([gx, gy, 1.25], dtype=np.float64)
beam_size = np.array([0.05, 1.25, 0.0], dtype=np.float64)
beam_rgba = np.array([0.22, 0.74, 0.97, 0.55], dtype=np.float32) # Glowing Cyan
mujoco.mjv_initGeom(dst_beam, mujoco.mjtGeom.mjGEOM_CYLINDER, beam_size, beam_pos, identity_mat, beam_rgba)
dst_beam.transparent = 1
scene.ngeom += 1
# 3. Ground target arrival ring
if scene.ngeom < scene.maxgeom:
dst_ring = scene.geoms[scene.ngeom]
ring_pos = np.array([gx, gy, 0.02], dtype=np.float64)
ring_size = np.array([self.goal_arrival_radius, 0.02, 0.0], dtype=np.float64)
ring_rgba = np.array([0.22, 0.74, 0.97, 0.45], dtype=np.float32)
mujoco.mjv_initGeom(dst_ring, mujoco.mjtGeom.mjGEOM_CYLINDER, ring_size, ring_pos, identity_mat, ring_rgba)
dst_ring.transparent = 1
scene.ngeom += 1
else:
for wp_idx, (wp_coord, wp_name, wp_radius) in enumerate(COURSE_WAYPOINTS):
if scene.ngeom + 3 >= scene.maxgeom:
break
wx = float(wp_coord[0])
wy = float(wp_coord[1])
wz = float(wp_coord[2])
if wp_idx == self.current_wp_idx:
# ACTIVE GOAL: Glowing floating target orb + vertical laser beacon column + ground target ring
# 1. Floating target orb (sphere)
dst_orb = scene.geoms[scene.ngeom]
orb_pos = np.array([wx, wy, wz + 0.35], dtype=np.float64)
orb_size = np.array([0.18, 0.18, 0.18], dtype=np.float64)
orb_rgba = np.array([0.98, 0.75, 0.15, 0.90], dtype=np.float32) # Vibrant Gold
mujoco.mjv_initGeom(dst_orb, mujoco.mjtGeom.mjGEOM_SPHERE, orb_size, orb_pos, identity_mat, orb_rgba)
dst_orb.transparent = 1
scene.ngeom += 1
# 2. Vertical laser beacon column (cylinder)
if scene.ngeom < scene.maxgeom:
dst_beam = scene.geoms[scene.ngeom]
beam_pos = np.array([wx, wy, wz + 0.60], dtype=np.float64)
beam_size = np.array([0.05, 0.60, 0.0], dtype=np.float64)
beam_rgba = np.array([0.22, 0.74, 0.97, 0.55], dtype=np.float32) # Glowing Cyan
mujoco.mjv_initGeom(dst_beam, mujoco.mjtGeom.mjGEOM_CYLINDER, beam_size, beam_pos, identity_mat, beam_rgba)
dst_beam.transparent = 1
scene.ngeom += 1
# 3. Ground target landing pad ring
if scene.ngeom < scene.maxgeom:
dst_ring = scene.geoms[scene.ngeom]
ring_pos = np.array([wx, wy, max(0.01, wz - 0.05)], dtype=np.float64)
ring_size = np.array([float(wp_radius) * 0.6, 0.02, 0.0], dtype=np.float64)
ring_rgba = np.array([0.22, 0.74, 0.97, 0.45], dtype=np.float32)
mujoco.mjv_initGeom(dst_ring, mujoco.mjtGeom.mjGEOM_CYLINDER, ring_size, ring_pos, identity_mat, ring_rgba)
dst_ring.transparent = 1
scene.ngeom += 1
elif wp_idx < self.current_wp_idx:
# COMPLETED GOAL: Subtle emerald disc on surface
dst_comp = scene.geoms[scene.ngeom]
comp_pos = np.array([wx, wy, max(0.01, wz - 0.05)], dtype=np.float64)
comp_size = np.array([0.25, 0.015, 0.0], dtype=np.float64)
comp_rgba = np.array([0.10, 0.85, 0.35, 0.40], dtype=np.float32) # Emerald Green
mujoco.mjv_initGeom(dst_comp, mujoco.mjtGeom.mjGEOM_CYLINDER, comp_size, comp_pos, identity_mat, comp_rgba)
dst_comp.transparent = 1
scene.ngeom += 1
else:
# UPCOMING GOAL: Translucent waypoint marker disc
dst_up = scene.geoms[scene.ngeom]
up_pos = np.array([wx, wy, max(0.01, wz - 0.05)], dtype=np.float64)
up_size = np.array([0.16, 0.015, 0.0], dtype=np.float64)
up_rgba = np.array([0.60, 0.65, 0.80, 0.25], dtype=np.float32) # Translucent Silver
mujoco.mjv_initGeom(dst_up, mujoco.mjtGeom.mjGEOM_CYLINDER, up_size, up_pos, identity_mat, up_rgba)
dst_up.transparent = 1
scene.ngeom += 1
# Inject translucent ghost avatars if provided
if ghost_datas:
scene = self._renderer.scene
for g_idx, g_data in enumerate(ghost_datas):
color = GHOST_PALETTE[g_idx % len(GHOST_PALETTE)]
if ghost_alpha is not None:
color = color.copy()
color[3] = float(ghost_alpha)
for gid in self.humanoid_geoms:
if scene.ngeom >= scene.maxgeom:
break
dst = scene.geoms[scene.ngeom]
pos = g_data.geom_xpos[gid].astype(np.float64)
size = self.model.geom_size[gid].astype(np.float64)
mat = g_data.geom_xmat[gid].astype(np.float64)
geom_type = int(self.model.geom_type[gid])
mujoco.mjv_initGeom(dst, geom_type, size, pos, mat, color)
dst.transparent = 1
scene.ngeom += 1
return self._renderer.render()