Upload 2 files

balance_plate_rl.py

@@ -0,0 +1,1666 @@
+#!/usr/bin/env python3
+"""
+PlateBalance RL: Advanced MuJoCo + PyTorch PPO Object Balancing System
+
+A high-performance MuJoCo + PyTorch PPO reinforcement learning system where
+a 2-axis tilting plate learns to balance 21 diverse rigid and multi-body objects:
+  - Standard Shapes: sphere, disk, egg, hollow cup, coin, stick, tall block,
+    triangular prism, cube block, puck.
+  - Harder Challenge Shapes: cone, horizontal capsule, ramp wedge, tetrahedron,
+    long flat bar, cross/plus, asymmetric L-shape, wide tile block, heavy bowling ball,
+    and off-center mass block.
+  - Multi-Body Crumbling Cookie: A breakable cookie with 5 independent physical crumb
+    fragments that scatter and slide independently, challenging the agent to keep
+    EVERY individual crumb on the plate!
+
+Installation:
+    pip install -U torch numpy mujoco imageio imageio-ffmpeg
+
+Usage:
+    # Train policy on all objects (default)
+    python balance_plate_rl.py
+
+    # Train or evaluate on a specific object (e.g., crumbling cookie or cone)
+    python balance_plate_rl.py --object cookie
+    python balance_plate_rl.py --object cone
+
+    # Resume training from latest checkpoint
+    python balance_plate_rl.py --resume
+
+    # Run full benchmark evaluation across all 21 objects
+    python balance_plate_rl.py --eval --checkpoint runs/plate_balance_v1/checkpoints/latest
+
+    # Interactive real-time 3D viewer (watch the cookie crumble and balance in real-time)
+    python balance_plate_rl.py --human-view --object cookie --checkpoint runs/plate_balance_v1/checkpoints/latest
+
+    # Record evaluation video
+    python balance_plate_rl.py --record-video --checkpoint runs/plate_balance_v1/checkpoints/latest
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import os
+import random
+import re
+import sys
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.distributions import Normal
+
+try:
+    import mujoco
+except ImportError as exc:
+    raise SystemExit("MuJoCo is required. Install with: pip install -U mujoco") from exc
+
+try:
+    import imageio.v2 as imageio
+except ImportError:
+    try:
+        import imageio
+    except ImportError:
+        imageio = None  # Optional for headless training/eval
+
+
+# =============================================================================
+# CONFIGURATION — DEFAULT VALUES
+# =============================================================================
+
+# ----------------------------- Experiment -----------------------------------
+SEED = 42
+EXPERIMENT_NAME = "plate_balance_v1"
+OUTPUT_DIR = "runs"
+RESUME = True
+RESUME_PATH = r"runs\plate_balance_v1\checkpoints\step_03825064_final"                 # Empty = automatically use latest checkpoint.
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# ----------------------------- Architecture ---------------------------------
+OBSERVATION_SIZE = 64
+HIDDEN_SIZE = 128                 # First/second hidden layer width.
+INTERMEDIATE_SIZE = 128           # Third hidden layer width.
+BOTTLENECK_SIZE = 64              # Fourth hidden layer width.
+ACTION_SIZE = 2                   # Roll torque, Pitch torque
+
+# Architecture: 64 -> 192 -> 192 -> 64 -> {actor: 2, critic: 1}
+ACTOR_LOG_STD_INIT = -0.75
+ORTHOGONAL_INIT = True
+
+# ----------------------------- PPO ------------------------------------------
+NUM_ENVS = 8
+ROLLOUT_STEPS = 512
+PPO_EPOCHS = 6
+MINIBATCH_SIZE = 2048
+GAMMA = 0.99
+GAE_LAMBDA = 0.95
+CLIP_COEF = 0.20
+VALUE_COEF = 0.50
+ENTROPY_COEF = 0.005
+MAX_GRAD_NORM = 0.50
+LEARNING_RATE = 3e-4
+ADAM_EPS = 1e-5
+ANNEAL_LR = True                  # Linearly anneal learning rate to 0.
+
+# ----------------------------- Training -------------------------------------
+NUM_EPISODES = 30_000
+MAX_EPISODE_STEPS = 10_000
+LOGGING_STEPS = 32240
+SAVE_STEPS = 1_000_000
+
+# ----------------------------- Video / Visualization ------------------------
+VISUALIZE = True                # True = save rollout videos periodically during training.
+VIDEO_EVERY_STEPS = 1_000_000
+VIDEO_LENGTH_STEPS = 3000
+VIDEO_FPS = 60
+RENDER_WIDTH = 640
+RENDER_HEIGHT = 480
+HUMAN_VIEW = False                # Optional interactive MuJoCo viewer.
+
+# ----------------------------- Physics & Environment ------------------------
+PHYSICS_TIMESTEP = 0.0025         # 400 Hz physics simulation.
+CONTROL_DECIMATION = 4            # 4 physics sub-steps per control step => 100 Hz RL control.
+GRAVITY = 9.81
+PLATE_HALF_SIZE = 0.90            # Half-width and half-length of square plate (meters).
+PLATE_THICKNESS = 0.08            # Half-thickness in z (box geom size z = 0.08).
+PLATE_FRICTION = (1.0, 0.02, 0.002) # Sliding, torsional, and rolling friction.
+PLATE_DAMPING = 0.08
+MAX_PLATE_ANGLE = math.radians(15.0)  # Maximum tilt angle in radians (~0.2618 rad).
+MAX_PLATE_TORQUE = 15.0           # Maximum plate torque in N*m.
+
+OBJECT_START_HEIGHT = 0.02        # Gentle initial clearance above plate when spawning (2cm).
+OBJECT_START_POS_RANGE = 0.35     # Spawn xy radius from plate center.
+OBJECT_START_VELOCITY = 0.15      # Initial linear velocity standard deviation.
+OBJECT_START_ANGULAR_VELOCITY = 0.20 # Initial angular velocity standard deviation.
+
+# Curriculum & Domain Randomization
+RANDOMIZE_MASS = True
+RANDOMIZE_FRICTION = True
+RANDOMIZE_OBJECT_SIZE = True
+MASS_LOG_RANGE = (0.50, 2.00)     # Multiplicative log-uniform mass range.
+FRICTION_RANGE = (0.45, 1.25)     # Multiplicative friction coefficient range.
+SIZE_RANGE = (0.85, 1.15)         # Multiplicative linear geometry scale range.
+RANDOMIZE_PLATE_ANGLE = True
+INITIAL_PLATE_ANGLE_RANGE = math.radians(3.0)
+
+# Reward Shaping
+CENTER_REWARD_SCALE = 2.0         # Reward for centering object on plate.
+VELOCITY_REWARD_SCALE = 0.40      # Reward for damping object linear velocity.
+ANGLE_PENALTY_SCALE = 0.20        # Penalty for excessive plate tilt.
+ACTION_PENALTY_SCALE = 0.005      # Penalty for excessive torque commands.
+ACTION_RATE_PENALTY_SCALE = 0.01  # Penalty for rapid actuator jerk (smooth control).
+EDGE_PENALTY_SCALE = 0.50         # Progressive penalty near plate perimeter.
+FALL_PENALTY = -10.0              # Terminal penalty when object falls off.
+SURVIVAL_REWARD = 0.05            # Step reward for keeping object on plate.
+
+# Fall detection margins
+FALL_MARGIN = 0.10                # Beyond plate half-size + margin triggers fall.
+FALL_Z_THRESHOLD = 0.90           # Height below which object is deemed fallen.
+
+# ----------------------------- Object Classification -----------------------
+STANDARD_OBJECT_TYPES = [
+    "sphere",
+    "disk",
+    "egg",
+    "cup",
+    "coin",
+    "stick",
+    "tall",
+    "triangle",
+    "block",
+    "puck",
+]
+
+HARDER_OBJECT_TYPES = [
+    "cone",
+    "capsule",
+    "wedge",
+    "tetra",
+    "flat_bar",
+    "cross",
+    "L_shape",
+    "wide_block",
+    "heavy_ball",
+    "offcenter_block",
+]
+
+MULTI_BODY_OBJECT_TYPES = [
+    "cookie",  # Multi-body crumbling cookie with 5 independent crumb fragments
+]
+OBJECT_TYPES = STANDARD_OBJECT_TYPES + HARDER_OBJECT_TYPES + MULTI_BODY_OBJECT_TYPES
+EVAL_OBJECT_TYPES = OBJECT_TYPES.copy()
+
+PRINT_OBJECT_COUNTS = True
+PRINT_HYPERPARAMS = True
+
+
+# =============================================================================
+# UTILITIES & REPRODUCIBILITY
+# =============================================================================
+
+def set_seed(seed: int) -> None:
+    """Set random seeds across Python, NumPy, and PyTorch for reproducibility."""
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+
+
+def extract_step_number(path: Path) -> int:
+    """Safely extract the integer step number from checkpoint folder names."""
+    matches = re.findall(r"\d+", path.name)
+    if matches:
+        return int(matches[0])
+    return -1
+
+
+def latest_checkpoint(root: Path) -> Optional[Path]:
+    """Find the latest checkpoint directory in a robust, crash-free manner."""
+    if not root.exists():
+        return None
+    latest_alias = root / "latest"
+    if latest_alias.exists() and (latest_alias / "model.pt").exists():
+        return latest_alias
+
+    candidates = [p for p in root.glob("step_*") if p.is_dir() and (p / "model.pt").exists()]
+    if not candidates:
+        return None
+    candidates.sort(key=extract_step_number)
+    return candidates[-1]
+
+
+def safe_float(x: Any) -> float:
+    return float(np.asarray(x).item())
+
+
+# =============================================================================
+# INERTIA CALCULATIONS & PHYSICAL DEFINITIONS
+# =============================================================================
+
+OBJECT_INFO: Dict[str, Dict[str, float]] = {
+    # Baseline 10 Shapes
+    "sphere": {"radius": 0.16, "mass": 0.55},
+    "disk": {"radius": 0.22, "height": 0.055, "mass": 0.60},
+    "egg": {"radius": 0.17, "height": 0.34, "mass": 0.58},
+    "cup": {"radius": 0.18, "height": 0.25, "mass": 0.48},
+    "coin": {"radius": 0.18, "height": 0.025, "mass": 0.25},
+    "stick": {"radius": 0.035, "length": 0.48, "mass": 0.22},
+    "tall": {"radius": 0.045, "height": 0.52, "mass": 0.35},
+    "triangle": {"size": 0.34, "height": 0.16, "mass": 0.50},
+    "block": {"size": 0.28, "mass": 0.70},
+    "puck": {"radius": 0.25, "height": 0.08, "mass": 0.70},
+
+    # Harder 10 Shapes
+    "cone": {"radius": 0.16, "height": 0.32, "mass": 0.50},
+    "capsule": {"radius": 0.065, "length": 0.36, "mass": 0.45},
+    "wedge": {"size": 0.36, "height": 0.20, "mass": 0.55},
+    "tetra": {"size": 0.36, "height": 0.30, "mass": 0.40},
+    "flat_bar": {"length": 0.76, "height": 0.04, "mass": 0.60},
+    "cross": {"size": 0.52, "height": 0.08, "mass": 0.65},
+    "L_shape": {"size": 0.32, "height": 0.08, "mass": 0.60},
+    "wide_block": {"size": 0.64, "height": 0.04, "mass": 0.75},
+    "heavy_ball": {"radius": 0.16, "mass": 3.50},
+    "offcenter_block": {"size": 0.36, "height": 0.20, "mass": 0.75},
+
+    # Multi-Body Crumbling Cookie Challenge (5 Crumb Fragments)
+    "cookie": {"radius": 0.22, "height": 0.04, "mass": 0.45},
+}
+
+
+def cube_inertia(m: float, sx: float, sy: float, sz: float) -> np.ndarray:
+    return np.array([
+        m * (sy * sy + sz * sz) / 12.0,
+        m * (sx * sx + sz * sz) / 12.0,
+        m * (sx * sx + sy * sy) / 12.0,
+    ], dtype=np.float64)
+
+
+def sphere_inertia(m: float, r: float) -> np.ndarray:
+    i = 0.4 * m * r * r
+    return np.array([i, i, i], dtype=np.float64)
+
+
+def cylinder_inertia(m: float, r: float, h: float) -> np.ndarray:
+    axial = 0.5 * m * r * r
+    transverse = m * (3.0 * r * r + h * h) / 12.0
+    return np.array([transverse, transverse, axial], dtype=np.float64)
+
+
+def capsule_inertia(m: float, r: float, length: float) -> np.ndarray:
+    rod = max(0.55 * m, 1e-6)
+    cyl = max(m - rod, 1e-6)
+    i_trans = rod * (length * length) / 12.0 + cyl * (3.0 * r * r + length * length) / 12.0
+    i_axial = 0.5 * cyl * r * r + rod * r * r / 2.0
+    return np.array([i_trans, i_trans, i_axial], dtype=np.float64)
+
+
+# Custom 3D Surface Meshes
+TRIANGLE_MESH = """
+<mesh name="triangle_mesh"
+      vertex="-0.30 -0.24 -0.08   0.30 -0.24 -0.08   0.00 0.30 -0.08
+               -0.30 -0.24  0.08   0.30 -0.24  0.08   0.00 0.30  0.08"
+      face="0 1 2   3 5 4   0 3 4   0 4 1   1 4 5   1 5 2   2 5 3   2 3 0" />
+"""
+
+CONE_MESH = """
+<mesh name="cone_mesh"
+      vertex=" 0.00  0.00  0.18
+               0.16  0.00 -0.14
+               0.11  0.11 -0.14
+               0.00  0.16 -0.14
+              -0.11  0.11 -0.14
+              -0.16  0.00 -0.14
+              -0.11 -0.11 -0.14
+               0.00 -0.16 -0.14
+               0.11 -0.11 -0.14
+               0.00  0.00 -0.14"
+      face="0 1 2   0 2 3   0 3 4   0 4 5   0 5 6   0 6 7   0 7 8   0 8 1
+            9 2 1   9 3 2   9 4 3   9 5 4   9 6 5   9 7 6   9 8 7   9 1 8" />
+"""
+
+WEDGE_MESH = """
+<mesh name="wedge_mesh"
+      vertex="-0.18 -0.15 -0.10
+               0.18 -0.15 -0.10
+               0.18  0.15 -0.10
+              -0.18  0.15 -0.10
+              -0.18 -0.15  0.10
+               0.18 -0.15  0.10"
+      face="0 1 2   0 2 3
+            0 4 5   0 5 1
+            1 5 2
+            0 3 4
+            3 2 5   3 5 4" />
+"""
+
+TETRA_MESH = """
+<mesh name="tetra_mesh"
+      vertex=" 0.00  0.00  0.18
+               0.18 -0.12 -0.12
+              -0.18 -0.12 -0.12
+               0.00  0.20 -0.12"
+      face="0 1 2   0 2 3   0 3 1   1 3 2" />
+"""
+
+
+def build_model_xml() -> str:
+    """Build and return the comprehensive MuJoCo XML containing all 21 objects and cookie crumbs."""
+    spawn_z = 1.05 + PLATE_THICKNESS + OBJECT_START_HEIGHT
+    return f"""
+<mujoco model="plate_balance">
+  <compiler angle="radian" coordinate="local" inertiafromgeom="auto" />
+  <option timestep="{PHYSICS_TIMESTEP}" gravity="0 0 -{GRAVITY}" integrator="implicitfast" />
+
+  <default>
+    <joint damping="{PLATE_DAMPING}" armature="0.01" />
+    <geom solref="0.008 1" solimp="0.90 0.95 0.01" />
+    <motor ctrllimited="true" ctrlrange="-1.0 1.0" />
+  </default>
+
+  <asset>
+    <texture type="skybox" builtin="gradient" rgb1="0.85 0.90 0.98" rgb2="0.65 0.75 0.90" width="512" height="512" />
+    <texture name="grid" type="2d" builtin="checker" width="512" height="512" rgb1="0.92 0.92 0.92" rgb2="0.80 0.80 0.80" />
+    <material name="grid_mat" texture="grid" texrepeat="5 5" reflectance="0.1" />
+    <texture name="cookie_tex" type="2d" builtin="checker" width="64" height="64" rgb1="0.82 0.58 0.32" rgb2="0.68 0.42 0.22" />
+    <material name="cookie_mat" texture="cookie_tex" reflectance="0.1" />
+    <material name="plate_mat" rgba="0.16 0.45 0.92 1" reflectance="0.3" specular="0.5" />
+    <material name="object_mat" rgba="0.95 0.45 0.16 1" reflectance="0.2" specular="0.4" />
+    <material name="heavy_mat" rgba="0.32 0.35 0.42 1" reflectance="0.6" specular="0.8" />
+    {TRIANGLE_MESH}
+    {CONE_MESH}
+    {WEDGE_MESH}
+    {TETRA_MESH}
+  </asset>
+
+  <worldbody>
+    <light directional="true" pos="0 -3 5" dir="0 0.5 -1" diffuse="0.8 0.8 0.8" specular="0.3 0.3 0.3" />
+    <light directional="true" pos="3 3 5" dir="-0.5 -0.5 -1" diffuse="0.4 0.4 0.4" />
+    <geom name="floor" type="plane" size="6 6 0.1" pos="0 0 0" material="grid_mat" contype="1" conaffinity="1" />
+
+    <!-- Cinematic cameras for visualization and video recording -->
+    <camera name="track" pos="0 -2.4 2.2" xyaxes="1 0 0  0 0.6 0.8" />
+    <camera name="isometric" pos="2.0 -2.0 2.3" xyaxes="0.707 0.707 0  -0.408 0.408 0.816" />
+    <camera name="top_down" pos="0 0 3.3" xyaxes="1 0 0  0 1 0" />
+    <camera name="side" pos="-2.8 0 1.5" xyaxes="0 -1 0  0.2 0 0.98" />
+
+    <!-- 2-Axis Gimbal Tilting Plate Mechanism -->
+    <body name="plate_x" pos="0 0 1.05">
+      <joint name="plate_roll" type="hinge" axis="1 0 0" range="-{MAX_PLATE_ANGLE} {MAX_PLATE_ANGLE}" limited="true" />
+      <inertial pos="0 0 0" mass="4.0" diaginertia="1.0 1.0 1.0" />
+
+      <body name="plate_y">
+        <joint name="plate_pitch" type="hinge" axis="0 1 0" range="-{MAX_PLATE_ANGLE} {MAX_PLATE_ANGLE}" limited="true" />
+        <inertial pos="0 0 0" mass="4.0" diaginertia="1.0 1.0 1.0" />
+        <geom name="plate" type="box" size="{PLATE_HALF_SIZE} {PLATE_HALF_SIZE} {PLATE_THICKNESS}"
+              material="plate_mat" friction="{PLATE_FRICTION[0]} {PLATE_FRICTION[1]} {PLATE_FRICTION[2]}"
+              contype="1" conaffinity="1" mass="4.0" />
+        <site name="plate_center" pos="0 0 {PLATE_THICKNESS + 0.005}" size="0.015" rgba="1 1 1 0.6" />
+      </body>
+    </body>
+
+    <!-- Primary candidate object body (for single rigid objects) -->
+    <body name="object" pos="0 0 {spawn_z}">
+      <freejoint name="object_free" />
+
+      <!-- Baseline 10 Shapes -->
+      <geom name="g_sphere" type="sphere" size="0.16" rgba="0.95 0.45 0.16 1" contype="1" conaffinity="1" />
+      <geom name="g_disk" type="cylinder" size="0.22 0.0275" rgba="0.95 0.45 0.16 1" contype="1" conaffinity="1" />
+      <geom name="g_egg" type="ellipsoid" size="0.135 0.135 0.19" rgba="0.95 0.45 0.16 1" contype="1" conaffinity="1" />
+
+      <!-- Hollow cup: base + 4 walls -->
+      <geom name="g_cup_bottom" type="cylinder" size="0.18 0.015" pos="0 0 -0.11" rgba="0.95 0.45 0.16 1" contype="1" conaffinity="1" />
+      <geom name="g_cup_w1" type="box" size="0.015 0.18 0.125" pos="0.165 0 0" rgba="0.95 0.45 0.16 1" contype="1" conaffinity="1" />
+      <geom name="g_cup_w2" type="box" size="0.015 0.18 0.125" pos="-0.165 0 0" rgba="0.95 0.45 0.16 1" contype="1" conaffinity="1" />
+      <geom name="g_cup_w3" type="box" size="0.15 0.015 0.125" pos="0 0.165 0" rgba="0.95 0.45 0.16 1" contype="1" conaffinity="1" />
+      <geom name="g_cup_w4" type="box" size="0.15 0.015 0.125" pos="0 -0.165 0" rgba="0.95 0.45 0.16 1" contype="1" conaffinity="1" />
+
+      <geom name="g_coin" type="cylinder" size="0.18 0.0125" rgba="0.95 0.45 0.16 1" contype="1" conaffinity="1" />
+      <geom name="g_stick" type="capsule" size="0.035 0.24" fromto="0 0 -0.24 0 0 0.24" rgba="0.95 0.45 0.16 1" contype="1" conaffinity="1" />
+      <geom name="g_tall" type="box" size="0.045 0.045 0.26" rgba="0.95 0.45 0.16 1" contype="1" conaffinity="1" />
+      <geom name="g_triangle" type="mesh" mesh="triangle_mesh" rgba="0.95 0.45 0.16 1" contype="1" conaffinity="1" />
+      <geom name="g_block" type="box" size="0.14 0.14 0.14" rgba="0.95 0.45 0.16 1" contype="1" conaffinity="1" />
+      <geom name="g_puck" type="cylinder" size="0.25 0.04" rgba="0.95 0.45 0.16 1" contype="1" conaffinity="1" />
+
+      <!-- Harder 10 Shapes -->
+      <geom name="g_cone" type="mesh" mesh="cone_mesh" rgba="0.95 0.45 0.16 1" contype="1" conaffinity="1" />
+      <geom name="g_capsule" type="capsule" size="0.065 0.18" fromto="-0.18 0 0  0.18 0 0" rgba="0.95 0.45 0.16 1" contype="1" conaffinity="1" />
+      <geom name="g_wedge" type="mesh" mesh="wedge_mesh" rgba="0.95 0.45 0.16 1" contype="1" conaffinity="1" />
+      <geom name="g_tetra" type="mesh" mesh="tetra_mesh" rgba="0.95 0.45 0.16 1" contype="1" conaffinity="1" />
+      <geom name="g_flat_bar" type="box" size="0.38 0.05 0.02" rgba="0.95 0.45 0.16 1" contype="1" conaffinity="1" />
+      <geom name="g_cross_1" type="box" size="0.26 0.055 0.04" rgba="0.95 0.45 0.16 1" contype="1" conaffinity="1" />
+      <geom name="g_cross_2" type="box" size="0.055 0.26 0.04" rgba="0.95 0.45 0.16 1" contype="1" conaffinity="1" />
+      <geom name="g_lshape_1" type="box" size="0.055 0.16 0.04" pos="0 -0.08 0" rgba="0.95 0.45 0.16 1" contype="1" conaffinity="1" />
+      <geom name="g_lshape_2" type="box" size="0.14 0.055 0.04" pos="0.085 0.08 0" rgba="0.95 0.45 0.16 1" contype="1" conaffinity="1" />
+      <geom name="g_wide_block" type="box" size="0.32 0.32 0.02" rgba="0.95 0.45 0.16 1" contype="1" conaffinity="1" />
+      <geom name="g_heavy_ball" type="sphere" size="0.16" material="heavy_mat" contype="1" conaffinity="1" />
+      <geom name="g_offcenter_block" type="box" size="0.18 0.18 0.10" rgba="0.95 0.45 0.16 1" contype="1" conaffinity="1" />
+    </body>
+
+    <!-- 5 Crumb bodies for the Crumbling Cookie challenge -->
+    <body name="crumb_0" pos="0 0 1.15">
+      <freejoint name="crumb_free_0" />
+      <geom name="g_crumb_0" type="cylinder" size="0.08 0.02" material="cookie_mat" contype="1" conaffinity="1" mass="0.15" />
+    </body>
+    <body name="crumb_1" pos="0.09 0 1.15">
+      <freejoint name="crumb_free_1" />
+      <geom name="g_crumb_1" type="box" size="0.035 0.035 0.02" material="cookie_mat" contype="1" conaffinity="1" mass="0.08" />
+    </body>
+    <body name="crumb_2" pos="-0.09 0 1.15">
+      <freejoint name="crumb_free_2" />
+      <geom name="g_crumb_2" type="cylinder" size="0.035 0.02" material="cookie_mat" contype="1" conaffinity="1" mass="0.08" />
+    </body>
+    <body name="crumb_3" pos="0 0.09 1.15">
+      <freejoint name="crumb_free_3" />
+      <geom name="g_crumb_3" type="box" size="0.04 0.03 0.02" material="cookie_mat" contype="1" conaffinity="1" mass="0.08" />
+    </body>
+    <body name="crumb_4" pos="0 -0.09 1.15">
+      <freejoint name="crumb_free_4" />
+      <geom name="g_crumb_4" type="sphere" size="0.03" material="cookie_mat" contype="1" conaffinity="1" mass="0.06" />
+    </body>
+  </worldbody>
+
+  <actuator>
+    <motor name="roll_motor" joint="plate_roll" gear="{MAX_PLATE_TORQUE}" />
+    <motor name="pitch_motor" joint="plate_pitch" gear="{MAX_PLATE_TORQUE}" />
+  </actuator>
+</mujoco>
+"""
+
+
+# =============================================================================
+# MUJOCO ENVIRONMENT (SINGLE & MULTI-BODY CRUMBLING COOKIE)
+# =============================================================================
+
+class PlateBalanceEnv:
+    """
+    Robust single-instance MuJoCo environment supporting 21 single & multi-body
+    objects with accurate physical scaling and multi-crumb tracking.
+    """
+
+    OBJECT_GEOM_NAMES = {
+        # Baseline 10
+        "sphere": ["g_sphere"],
+        "disk": ["g_disk"],
+        "egg": ["g_egg"],
+        "cup": ["g_cup_bottom", "g_cup_w1", "g_cup_w2", "g_cup_w3", "g_cup_w4"],
+        "coin": ["g_coin"],
+        "stick": ["g_stick"],
+        "tall": ["g_tall"],
+        "triangle": ["g_triangle"],
+        "block": ["g_block"],
+        "puck": ["g_puck"],
+        # Harder 10
+        "cone": ["g_cone"],
+        "capsule": ["g_capsule"],
+        "wedge": ["g_wedge"],
+        "tetra": ["g_tetra"],
+        "flat_bar": ["g_flat_bar"],
+        "cross": ["g_cross_1", "g_cross_2"],
+        "L_shape": ["g_lshape_1", "g_lshape_2"],
+        "wide_block": ["g_wide_block"],
+        "heavy_ball": ["g_heavy_ball"],
+        "offcenter_block": ["g_offcenter_block"],
+    }
+
+    def __init__(self, seed: int, render: bool = False, active_objects: Optional[List[str]] = None):
+        self.rng = np.random.default_rng(seed)
+        self.model = mujoco.MjModel.from_xml_string(build_model_xml())
+        self.data = mujoco.MjData(self.model)
+        self.model.opt.timestep = PHYSICS_TIMESTEP
+        self.render_enabled = render
+        self.renderer: Optional[mujoco.Renderer] = None
+        self.active_objects = list(active_objects) if active_objects else list(OBJECT_TYPES)
+
+        if render:
+            self.renderer = mujoco.Renderer(self.model, height=RENDER_HEIGHT, width=RENDER_WIDTH)
+
+        # Primary Single Object
+        self.qpos_object = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "object_free")
+        self.object_qpos_addr = self.model.jnt_qposadr[self.qpos_object]
+        self.object_dof_addr = self.model.jnt_dofadr[self.qpos_object]
+        self.object_body_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_BODY, "object")
+
+        # Plate Joints and Sites
+        self.plate_roll_joint = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "plate_roll")
+        self.plate_pitch_joint = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, "plate_pitch")
+        self.plate_roll_qpos = self.model.jnt_qposadr[self.plate_roll_joint]
+        self.plate_pitch_qpos = self.model.jnt_qposadr[self.plate_pitch_joint]
+        self.plate_roll_dof = self.model.jnt_dofadr[self.plate_roll_joint]
+        self.plate_pitch_dof = self.model.jnt_dofadr[self.plate_pitch_joint]
+        self.plate_body_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_BODY, "plate_y")
+        self.plate_center_site = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_SITE, "plate_center")
+        self.plate_geom_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, "plate")
+
+        # Cookie Crumb Bodies (5 fragments)
+        self.crumb_body_ids = [mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_BODY, f"crumb_{i}") for i in range(5)]
+        self.crumb_qpos_addrs = [self.model.jnt_qposadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, f"crumb_free_{i}")] for i in range(5)]
+        self.crumb_dof_addrs = [self.model.jnt_dofadr[mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, f"crumb_free_{i}")] for i in range(5)]
+        self.crumb_geom_ids = [mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, f"g_crumb_{i}") for i in range(5)]
+
+        # Map candidate geoms
+        self.object_geom_ids = {
+            name: [mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_GEOM, g) for g in geoms]
+            for name, geoms in self.OBJECT_GEOM_NAMES.items()
+        }
+
+        # Cache baseline sizes and positions for physical scaling
+        self.base_geom_sizes: Dict[int, np.ndarray] = {}
+        self.base_geom_positions: Dict[int, np.ndarray] = {}
+        for ids in self.object_geom_ids.values():
+            for gid in ids:
+                self.base_geom_sizes[gid] = self.model.geom_size[gid].copy()
+                self.base_geom_positions[gid] = self.model.geom_pos[gid].copy()
+
+        for gid in self.crumb_geom_ids:
+            self.base_geom_sizes[gid] = self.model.geom_size[gid].copy()
+            self.base_geom_positions[gid] = self.model.geom_pos[gid].copy()
+
+        self.steps = 0
+        self.episode_reward = 0.0
+        self.current_object = "sphere"
+        self.object_mass = 0.55
+        self.object_scale = 1.0
+        self.base_friction = 0.90
+        self.last_action = np.zeros(ACTION_SIZE, dtype=np.float32)
+
+    def _set_object_geometry(self, object_name: str) -> None:
+        """Activate the selected single or multi-body object and scale its physics."""
+        # Deactivate all single object geoms
+        for ids in self.object_geom_ids.values():
+            for gid in ids:
+                self.model.geom_size[gid, :] = 1e-6
+                self.model.geom_pos[gid, :] = np.array([0.0, 0.0, -999.0])
+                self.model.geom_rgba[gid, 3] = 0.0
+
+        # Deactivate all crumb geoms by default
+        for gid in self.crumb_geom_ids:
+            self.model.geom_size[gid, :] = 1e-6
+            self.model.geom_pos[gid, :] = np.array([0.0, 0.0, -999.0])
+            self.model.geom_rgba[gid, 3] = 0.0
+
+        s = self.object_scale
+        info = OBJECT_INFO[object_name]
+        m = self.object_mass
+
+        if object_name == "cookie":
+            # Activate 5 cookie crumb fragments
+            for i, gid in enumerate(self.crumb_geom_ids):
+                self.model.geom_friction[gid, 0] = self.base_friction
+                self.model.geom_rgba[gid, :] = np.array([0.82, 0.58, 0.32, 1.0])
+                self.model.geom_size[gid, :] = self.base_geom_sizes[gid] * s
+                self.model.geom_pos[gid, :] = self.base_geom_positions[gid] * s
+                self.model.body_mass[self.crumb_body_ids[i]] = (m / 5.0)
+            return
+
+        # Activate single candidate object
+        active = self.object_geom_ids[object_name]
+        rgba = np.array([0.35, 0.38, 0.45, 1.0]) if object_name == "heavy_ball" else np.array([0.95, 0.45, 0.16, 1.0])
+        for gid in active:
+            self.model.geom_friction[gid, 0] = self.base_friction
+            self.model.geom_rgba[gid, :] = rgba
+            self.model.geom_size[gid, :] = self.base_geom_sizes[gid] * s
+            self.model.geom_pos[gid, :] = self.base_geom_positions[gid] * s
+
+        # Reset body center of mass position (offset for offcenter_block)
+        if object_name == "offcenter_block":
+            self.model.body_ipos[self.object_body_id] = np.array([0.08 * s, 0.06 * s, -0.02 * s])
+        else:
+            self.model.body_ipos[self.object_body_id] = np.array([0.0, 0.0, 0.0])
+
+        # Compute accurate 3D moment of inertia tensor
+        if object_name in ("sphere", "heavy_ball"):
+            inertia = sphere_inertia(m, info["radius"] * s)
+        elif object_name in ("disk", "coin", "puck"):
+            inertia = cylinder_inertia(m, info["radius"] * s, info["height"] * s)
+        elif object_name == "egg":
+            rx = info["radius"] * 0.85 * s
+            rz = info["height"] * 0.55 * s
+            inertia = np.array([
+                m * (rx * rx + rz * rz) / 5.0,
+                m * (rx * rx + rz * rz) / 5.0,
+                m * (2.0 * rx * rx) / 5.0,
+            ])
+        elif object_name == "cup":
+            inertia = cylinder_inertia(m, info["radius"] * s, info["height"] * s)
+        elif object_name in ("stick", "capsule"):
+            inertia = capsule_inertia(m, info["radius"] * s, info["length"] * s)
+        elif object_name == "tall":
+            side = 0.09 * s
+            height = info["height"] * s
+            inertia = cube_inertia(m, side, side, height)
+        elif object_name == "cone":
+            r = info["radius"] * s
+            h = info["height"] * s
+            i_trans = 0.6 * m * (0.25 * r * r + h * h)
+            i_axial = 0.3 * m * r * r
+            inertia = np.array([i_trans, i_trans, i_axial])
+        elif object_name in ("triangle", "wedge"):
+            side = info["size"] * s
+            height = info["height"] * s
+            inertia = cube_inertia(m, side, side, height) * 1.10
+        elif object_name == "tetra":
+            side = info["size"] * s
+            i_val = (m * side * side) / 20.0
+            inertia = np.array([i_val, i_val, i_val])
+        elif object_name == "flat_bar":
+            inertia = cube_inertia(m, info["length"] * s, 0.10 * s, info["height"] * s)
+        elif object_name == "cross":
+            side = info["size"] * s
+            h = info["height"] * s
+            inertia = cube_inertia(m, side, side, h) * 0.70
+        elif object_name == "L_shape":
+            side = info["size"] * s
+            h = info["height"] * s
+            inertia = cube_inertia(m, side, side, h) * 0.85
+        elif object_name == "wide_block":
+            side = info["size"] * s
+            h = info["height"] * s
+            inertia = cube_inertia(m, side, side, h)
+        elif object_name == "offcenter_block":
+            side = info["size"] * s
+            h = info["height"] * s
+            inertia = cube_inertia(m, side, side, h) * 1.20
+        else:  # block
+            side = info["size"] * s
+            inertia = cube_inertia(m, side, side, side)
+
+        self.model.body_mass[self.object_body_id] = m
+        self.model.body_inertia[self.object_body_id, :] = np.maximum(inertia, 1e-6)
+
+    def reset(self, specific_object: Optional[str] = None) -> np.ndarray:
+        """Reset the environment state for a new episode."""
+        if specific_object and specific_object in OBJECT_INFO:
+            self.current_object = specific_object
+        else:
+            self.current_object = str(self.rng.choice(self.active_objects))
+
+        info = OBJECT_INFO[self.current_object]
+        self.object_mass = info["mass"]
+        if RANDOMIZE_MASS:
+            mult = float(np.exp(self.rng.uniform(math.log(MASS_LOG_RANGE[0]), math.log(MASS_LOG_RANGE[1]))))
+            self.object_mass *= mult
+
+        self.object_scale = float(self.rng.uniform(*SIZE_RANGE)) if RANDOMIZE_OBJECT_SIZE else 1.0
+        self.base_friction = float(self.rng.uniform(*FRICTION_RANGE)) if RANDOMIZE_FRICTION else 0.90
+        self._set_object_geometry(self.current_object)
+
+        mujoco.mj_resetData(self.model, self.data)
+
+        # Initial plate tilt randomization
+        if RANDOMIZE_PLATE_ANGLE:
+            self.data.qpos[self.plate_roll_qpos] = self.rng.uniform(-INITIAL_PLATE_ANGLE_RANGE, INITIAL_PLATE_ANGLE_RANGE)
+            self.data.qpos[self.plate_pitch_qpos] = self.rng.uniform(-INITIAL_PLATE_ANGLE_RANGE, INITIAL_PLATE_ANGLE_RANGE)
+        else:
+            self.data.qpos[self.plate_roll_qpos] = 0.0
+            self.data.qpos[self.plate_pitch_qpos] = 0.0
+
+        x = float(self.rng.uniform(-OBJECT_START_POS_RANGE, OBJECT_START_POS_RANGE))
+        y = float(self.rng.uniform(-OBJECT_START_POS_RANGE, OBJECT_START_POS_RANGE))
+
+        if self.current_object == "cookie":
+            # Park single object far away
+            base_s = self.object_qpos_addr
+            self.data.qpos[base_s:base_s + 7] = np.array([0.0, 0.0, -999.0, 1.0, 0.0, 0.0, 0.0])
+            self.data.qvel[self.object_dof_addr:self.object_dof_addr + 6] = 0.0
+
+            # Spawn 5 cookie crumb fragments on the plate with scatter offsets & velocities
+            crumb_offsets = [
+                (0.0, 0.0),
+                (0.08 * self.object_scale, 0.0),
+                (-0.08 * self.object_scale, 0.0),
+                (0.0, 0.08 * self.object_scale),
+                (0.0, -0.08 * self.object_scale),
+            ]
+            for i in range(5):
+                dx, dy = crumb_offsets[i]
+                dx += float(self.rng.uniform(-0.02, 0.02))
+                dy += float(self.rng.uniform(-0.02, 0.02))
+                c_addr = self.crumb_qpos_addrs[i]
+                c_dof = self.crumb_dof_addrs[i]
+                z_c = 1.05 + PLATE_THICKNESS + 0.02 + OBJECT_START_HEIGHT
+                self.data.qpos[c_addr:c_addr + 7] = np.array([x + dx, y + dy, z_c, 1.0, 0.0, 0.0, 0.0])
+                self.data.qvel[c_dof:c_dof + 6] = np.concatenate([
+                    self.rng.normal(0.0, OBJECT_START_VELOCITY, 3),
+                    self.rng.normal(0.0, OBJECT_START_ANGULAR_VELOCITY, 3),
+                ])
+        else:
+            # Park all crumb fragments far away
+            for i in range(5):
+                c_addr = self.crumb_qpos_addrs[i]
+                c_dof = self.crumb_dof_addrs[i]
+                self.data.qpos[c_addr:c_addr + 7] = np.array([0.0, 0.0, -999.0, 1.0, 0.0, 0.0, 0.0])
+                self.data.qvel[c_dof:c_dof + 6] = 0.0
+
+            # Spawn single object with safe clearance
+            base = self.object_qpos_addr
+            if self.current_object in ("sphere", "heavy_ball"):
+                half_height = info["radius"] * self.object_scale
+            elif self.current_object in ("stick", "capsule", "flat_bar"):
+                half_height = (info["length"] / 2.0) * self.object_scale if "length" in info else (info["height"] / 2.0) * self.object_scale
+            elif self.current_object in ("block", "cross", "L_shape", "wide_block", "wedge", "tetra", "offcenter_block"):
+                half_height = (info.get("height", info.get("size", 0.20)) / 2.0) * self.object_scale
+            else:
+                half_height = (info["height"] / 2.0) * self.object_scale
+
+            z = 1.05 + PLATE_THICKNESS + half_height + OBJECT_START_HEIGHT
+            self.data.qpos[base:base + 7] = np.array([x, y, z, 1.0, 0.0, 0.0, 0.0])
+            self.data.qvel[self.object_dof_addr:self.object_dof_addr + 6] = np.concatenate([
+                self.rng.normal(0.0, OBJECT_START_VELOCITY, 3),
+                self.rng.normal(0.0, OBJECT_START_ANGULAR_VELOCITY, 3),
+            ])
+
+        self.data.ctrl[:] = 0.0
+        mujoco.mj_forward(self.model, self.data)
+        self.steps = 0
+        self.episode_reward = 0.0
+        self.last_action = np.zeros(ACTION_SIZE, dtype=np.float32)
+        return self._observation()
+
+    def _observation(self) -> np.ndarray:
+        """Construct the normalized 64-dimensional observation vector."""
+        plate_angles = np.array([self.data.qpos[self.plate_roll_qpos], self.data.qpos[self.plate_pitch_qpos]], dtype=np.float32)
+        plate_vel = np.array([self.data.qvel[self.plate_roll_dof], self.data.qvel[self.plate_pitch_dof]], dtype=np.float32)
+        plate_surface_center = self.data.site_xpos[self.plate_center_site]
+
+        info = OBJECT_INFO[self.current_object]
+        object_descriptor = np.array([
+            self.object_mass,
+            self.object_scale,
+            self.base_friction,
+            info.get("radius", info.get("size", 0.10)),
+            info.get("height", info.get("length", 0.10)),
+        ], dtype=np.float32)
+
+        if self.current_object == "cookie":
+            # Multi-body crumb tracking
+            crumb_xys = [self.data.xpos[bid, :2] for bid in self.crumb_body_ids]
+            crumb_zs = [self.data.qpos[addr + 2] for addr in self.crumb_qpos_addrs]
+            on_plate_crumbs = [
+                p for p, z in zip(crumb_xys, crumb_zs)
+                if max(abs(p[0]), abs(p[1])) <= (PLATE_HALF_SIZE + FALL_MARGIN) and z >= FALL_Z_THRESHOLD
+            ]
+            if on_plate_crumbs:
+                centroid_xy = np.mean(on_plate_crumbs, axis=0)
+                mean_z = float(np.mean([z for z in crumb_zs if z >= FALL_Z_THRESHOLD]))
+                spread = float(np.max([np.linalg.norm(p - centroid_xy) for p in on_plate_crumbs]))
+            else:
+                centroid_xy = np.zeros(2, dtype=np.float32)
+                mean_z = 1.15
+                spread = 0.0
+
+            object_pos = np.array([centroid_xy[0], centroid_xy[1], mean_z], dtype=np.float32)
+            rel_xy = (centroid_xy - plate_surface_center[:2]).astype(np.float32)
+            object_linvel = np.mean([self.data.qvel[dof:dof+3] for dof in self.crumb_dof_addrs], axis=0).astype(np.float32)
+            object_angvel = np.zeros(3, dtype=np.float32)
+            object_quat = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32)
+
+            # Crumb specific relative features
+            crumb_ratio = float(len(on_plate_crumbs) / 5.0)
+            crumb_rel_1 = (crumb_xys[1] - centroid_xy).astype(np.float32)
+            crumb_rel_2 = (crumb_xys[2] - centroid_xy).astype(np.float32)
+            multi_body_feats = np.concatenate([
+                np.array([1.0, spread, crumb_ratio], dtype=np.float32),
+                crumb_rel_1,
+                crumb_rel_2,
+            ])
+        else:
+            base_q = self.object_qpos_addr
+            base_v = self.object_dof_addr
+            object_pos = self.data.qpos[base_q:base_q + 3].astype(np.float32)
+            object_quat = self.data.qpos[base_q + 3:base_q + 7].astype(np.float32)
+            object_linvel = self.data.qvel[base_v:base_v + 3].astype(np.float32)
+            object_angvel = self.data.qvel[base_v + 3:base_v + 6].astype(np.float32)
+            rel_pos = object_pos - plate_surface_center
+            rel_xy = rel_pos[:2].astype(np.float32)
+            multi_body_feats = np.zeros(7, dtype=np.float32)
+
+        state = np.concatenate([
+            object_pos,              # 3
+            rel_xy,                  # 2
+            object_linvel,           # 3
+            object_angvel,           # 3
+            object_quat,             # 4
+            plate_angles,            # 2
+            plate_vel,               # 2
+            self.data.qfrc_actuator[:2], # 2
+            self.last_action,        # 2
+            object_descriptor,       # 5
+            multi_body_feats,        # 7
+        ]).astype(np.float32)
+
+        out = np.zeros(OBSERVATION_SIZE, dtype=np.float32)
+        n = min(len(state), OBSERVATION_SIZE)
+        out[:n] = state[:n]
+        out = np.clip(out, -10.0, 10.0)
+        return out
+
+    def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, bool, Dict[str, Any]]:
+        """Advance simulation and compute reward for single or multi-crumb cookie objects."""
+        action = np.asarray(action, dtype=np.float64)
+        action = np.clip(action, -1.0, 1.0)
+        self.data.ctrl[0] = action[0]
+        self.data.ctrl[1] = action[1]
+
+        for _ in range(CONTROL_DECIMATION):
+            mujoco.mj_step(self.model, self.data)
+
+        self.steps += 1
+        plate_ang = np.array([self.data.qpos[self.plate_roll_qpos], self.data.qpos[self.plate_pitch_qpos]])
+        plate_ang_mag = float(np.linalg.norm(plate_ang))
+        action_rate_penalty = float(np.mean((action - self.last_action) ** 2))
+        self.last_action = action.astype(np.float32).copy()
+
+        if self.current_object == "cookie":
+            # Multi-body Crumbling Cookie Evaluation
+            crumb_xys = [self.data.xpos[bid, :2] for bid in self.crumb_body_ids]
+            crumb_zs = [self.data.qpos[addr + 2] for addr in self.crumb_qpos_addrs]
+
+            on_plate_flags = [
+                bool(max(abs(p[0]), abs(p[1])) <= (PLATE_HALF_SIZE + FALL_MARGIN) and z >= FALL_Z_THRESHOLD)
+                for p, z in zip(crumb_xys, crumb_zs)
+            ]
+            crumbs_on = sum(on_plate_flags)
+            crumb_ratio = crumbs_on / 5.0
+
+            dists = [float(np.linalg.norm(p)) for p in crumb_xys]
+            max_dist = max(dists)
+            mean_dist = float(np.mean(dists))
+            box_dist = float(max(max(abs(p[0]), abs(p[1])) for p in crumb_xys))
+
+            vels = [float(np.linalg.norm(self.data.qvel[dof:dof+2])) for dof in self.crumb_dof_addrs]
+            mean_vel = float(np.mean(vels))
+
+            center_score = 0.5 * math.exp(-4.0 * mean_dist * mean_dist) + 0.5 * math.exp(-4.0 * max_dist * max_dist)
+            velocity_score = math.exp(-1.8 * mean_vel * mean_vel)
+            edge_fraction = np.clip(box_dist / PLATE_HALF_SIZE, 0.0, 1.5)
+            edge_penalty = max(0.0, float(edge_fraction) - 0.65) ** 2
+
+            # Reward scaled by the fraction of cookie crumbs kept on the plate
+            reward = (
+                SURVIVAL_REWARD * crumb_ratio
+                + CENTER_REWARD_SCALE * center_score
+                + VELOCITY_REWARD_SCALE * velocity_score
+                - ANGLE_PENALTY_SCALE * (plate_ang_mag / MAX_PLATE_ANGLE) ** 2
+                - ACTION_PENALTY_SCALE * float(np.mean(action ** 2))
+                - ACTION_RATE_PENALTY_SCALE * action_rate_penalty
+                - EDGE_PENALTY_SCALE * edge_penalty
+            )
+
+            # Failure occurs when any crumb falls off (or partial penalty for lost crumbs)
+            fallen = bool(crumbs_on < 5)
+            dist_xy = mean_dist
+            vel_xy = mean_vel
+        else:
+            # Single Rigid Object Evaluation
+            pos = self.data.xpos[self.object_body_id, :2]
+            dist_xy = float(np.linalg.norm(pos))
+            box_dist = float(max(abs(pos[0]), abs(pos[1])))
+
+            obj_vel = self.data.qvel[self.object_dof_addr:self.object_dof_addr + 3]
+            vel_xy = float(np.linalg.norm(obj_vel[:2]))
+
+            center_score = math.exp(-4.0 * dist_xy * dist_xy)
+            velocity_score = math.exp(-1.8 * vel_xy * vel_xy)
+            edge_fraction = np.clip(box_dist / PLATE_HALF_SIZE, 0.0, 1.5)
+            edge_penalty = max(0.0, float(edge_fraction) - 0.65) ** 2
+
+            reward = (
+                SURVIVAL_REWARD
+                + CENTER_REWARD_SCALE * center_score
+                + VELOCITY_REWARD_SCALE * velocity_score
+                - ANGLE_PENALTY_SCALE * (plate_ang_mag / MAX_PLATE_ANGLE) ** 2
+                - ACTION_PENALTY_SCALE * float(np.mean(action ** 2))
+                - ACTION_RATE_PENALTY_SCALE * action_rate_penalty
+                - EDGE_PENALTY_SCALE * edge_penalty
+            )
+
+            object_z = self.data.qpos[self.object_qpos_addr + 2]
+            fallen = bool(box_dist > (PLATE_HALF_SIZE + FALL_MARGIN) or object_z < FALL_Z_THRESHOLD)
+            crumbs_on = 1 if not fallen else 0
+
+        timeout = bool(self.steps >= MAX_EPISODE_STEPS)
+        terminated = fallen
+        truncated = timeout and not fallen
+
+        if fallen:
+            reward += FALL_PENALTY
+
+        self.episode_reward += reward
+        info = {
+            "object": self.current_object,
+            "distance": dist_xy,
+            "box_distance": box_dist,
+            "object_velocity": vel_xy,
+            "fallen": fallen,
+            "timeout": timeout,
+            "crumbs_on_plate": crumbs_on,
+            "episode_reward": self.episode_reward,
+            "episode_length": self.steps,
+        }
+        return self._observation(), float(reward), terminated, truncated, info
+
+    def render(self, camera: str = "track") -> np.ndarray:
+        """Render RGB visual frame from the specified camera."""
+        if not self.render_enabled:
+            raise RuntimeError("Environment initialized with render=False")
+        assert self.renderer is not None
+        self.renderer.update_scene(self.data, camera=camera)
+        return self.renderer.render().copy()
+
+    def close(self) -> None:
+        """Cleanly release rendering and physics resources."""
+        if self.renderer is not None:
+            self.renderer.close()
+            self.renderer = None
+
+
+# =============================================================================
+# ACTOR / CRITIC NEURAL NETWORK
+# =============================================================================
+
+class ActorCritic(nn.Module):
+    """
+    Continuous-action Actor-Critic MLP with Tanh-squashed Gaussian policy
+    and exact, numerically stable log-probability computation.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.backbone = nn.Sequential(
+            nn.Linear(OBSERVATION_SIZE, HIDDEN_SIZE),
+            nn.SiLU(),
+            nn.Linear(HIDDEN_SIZE, INTERMEDIATE_SIZE),
+            nn.SiLU(),
+            nn.Linear(INTERMEDIATE_SIZE, BOTTLENECK_SIZE),
+            nn.SiLU(),
+        )
+        self.actor = nn.Linear(BOTTLENECK_SIZE, ACTION_SIZE)
+        self.critic = nn.Linear(BOTTLENECK_SIZE, 1)
+        self.log_std = nn.Parameter(torch.full((ACTION_SIZE,), ACTOR_LOG_STD_INIT))
+        self._init_weights()
+
+    def _init_weights(self) -> None:
+        if not ORTHOGONAL_INIT:
+            return
+        for layer in self.backbone:
+            if isinstance(layer, nn.Linear):
+                nn.init.orthogonal_(layer.weight, gain=math.sqrt(2.0))
+                nn.init.zeros_(layer.bias)
+        nn.init.orthogonal_(self.actor.weight, gain=0.01)
+        nn.init.zeros_(self.actor.bias)
+        nn.init.orthogonal_(self.critic.weight, gain=1.0)
+        nn.init.zeros_(self.critic.bias)
+
+    def forward(self, obs: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        h = self.backbone(obs)
+        mean = self.actor(h)
+        value = self.critic(h).squeeze(-1)
+        std = self.log_std.exp().expand_as(mean)
+        return mean, std, value
+
+    def get_value(self, obs: torch.Tensor) -> torch.Tensor:
+        return self.critic(self.backbone(obs)).squeeze(-1)
+
+    def get_action_and_value(
+        self, obs: torch.Tensor, raw_action: Optional[torch.Tensor] = None, deterministic: bool = False
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        mean, std, value = self(obs)
+        dist = Normal(mean, std)
+
+        if raw_action is None:
+            raw_action = mean if deterministic else dist.sample()
+
+        squashed = torch.tanh(raw_action)
+        # Numerically stable exact log-prob correction for Tanh transformation
+        log_prob = dist.log_prob(raw_action).sum(-1)
+        log_prob -= torch.log(torch.clamp(1.0 - squashed.pow(2), min=1e-6)).sum(-1)
+        entropy = dist.entropy().sum(-1)
+        return squashed, log_prob, entropy, value, raw_action
+
+
+# =============================================================================
+# PPO TRAINER
+# =============================================================================
+
+@dataclass
+class EpisodeStats:
+    reward: float = 0.0
+    length: int = 0
+    object_name: str = ""
+    success: bool = False
+    crumbs_on: int = 1
+
+
+class PPOTrainer:
+    """High-throughput PPO Trainer with robust GAE, metric logging, and checkpointing."""
+
+    def __init__(self, device: Optional[str] = None):
+        self.device = torch.device(device or DEVICE)
+        self.policy = ActorCritic().to(self.device)
+        self.optimizer = torch.optim.Adam(self.policy.parameters(), lr=LEARNING_RATE, eps=ADAM_EPS)
+        self.global_steps = 0
+        self.episodes = 0
+        self.updates = 0
+        self.last_log_step = 0
+        self.last_save_step = 0
+        self.last_video_step = 0
+        self.object_counts: Dict[str, int] = {name: 0 for name in OBJECT_TYPES}
+
+    def save(self, directory: Path, latest: bool = True, extra: Optional[Dict[str, Any]] = None) -> Path:
+        """Save a complete resumable checkpoint."""
+        directory.mkdir(parents=True, exist_ok=True)
+        model_path = directory / "model.pt"
+        state_path = directory / "trainer_state.pt"
+        torch.save(self.policy.state_dict(), model_path)
+        trainer_state = {
+            "optimizer": self.optimizer.state_dict(),
+            "global_steps": self.global_steps,
+            "episodes": self.episodes,
+            "updates": self.updates,
+            "last_log_step": self.last_log_step,
+            "last_save_step": self.last_save_step,
+            "last_video_step": self.last_video_step,
+            "object_counts": self.object_counts,
+            "torch_rng_state": torch.get_rng_state(),
+            "numpy_rng_state": np.random.get_state(),
+            "python_rng_state": random.getstate(),
+            "config": {
+                k: v for k, v in globals().items()
+                if k.isupper() and isinstance(v, (int, float, str, bool, tuple, list))
+            },
+            "extra": extra or {},
+        }
+        if torch.cuda.is_available():
+            trainer_state["cuda_rng_state"] = torch.cuda.get_rng_state_all()
+        torch.save(trainer_state, state_path)
+
+        if latest:
+            latest_dir = directory.parent / "latest"
+            latest_dir.mkdir(parents=True, exist_ok=True)
+            torch.save(self.policy.state_dict(), latest_dir / "model.pt")
+            torch.save(trainer_state, latest_dir / "trainer_state.pt")
+            with open(latest_dir / "config.json", "w", encoding="utf-8") as f:
+                json.dump(trainer_state["config"], f, indent=2, default=str)
+        return directory
+
+    def load(self, directory: Path) -> None:
+        """Restore policy weights, optimizer state, and RNGs from checkpoint."""
+        model_path = directory / "model.pt"
+        state_path = directory / "trainer_state.pt"
+        if not model_path.exists() or not state_path.exists():
+            raise FileNotFoundError(f"Checkpoint directory missing model.pt or trainer_state.pt: {directory}")
+
+        self.policy.load_state_dict(torch.load(model_path, map_location=self.device, weights_only=True))
+        state = torch.load(state_path, map_location="cpu", weights_only=False)
+        self.optimizer.load_state_dict(state["optimizer"])
+        self.global_steps = int(state.get("global_steps", 0))
+        self.episodes = int(state.get("episodes", 0))
+        self.updates = int(state.get("updates", 0))
+        self.last_log_step = int(state.get("last_log_step", 0))
+        self.last_save_step = int(state.get("last_save_step", 0))
+        self.last_video_step = int(state.get("last_video_step", 0))
+        self.object_counts.update(state.get("object_counts", {}))
+
+        try:
+            torch.set_rng_state(state["torch_rng_state"])
+            np.random.set_state(state["numpy_rng_state"])
+            random.setstate(state["python_rng_state"])
+            if torch.cuda.is_available() and "cuda_rng_state" in state:
+                torch.cuda.set_rng_state_all(state["cuda_rng_state"])
+        except Exception as exc:
+            print(f"[resume] Notice: Could not restore full RNG state: {exc}")
+
+        print(f"[resume] Loaded checkpoint {directory} | global_steps={self.global_steps:,} episodes={self.episodes:,}")
+
+    def update(self, batch: Dict[str, np.ndarray], total_steps_target: int = NUM_EPISODES * 100) -> Dict[str, float]:
+        """Perform PPO mini-batch updates over the collected rollout."""
+        obs = torch.as_tensor(batch["obs"], dtype=torch.float32, device=self.device)
+        actions = torch.as_tensor(batch["raw_actions"], dtype=torch.float32, device=self.device)
+        old_logprobs = torch.as_tensor(batch["logprobs"], dtype=torch.float32, device=self.device)
+        advantages = torch.as_tensor(batch["advantages"], dtype=torch.float32, device=self.device)
+        returns = torch.as_tensor(batch["returns"], dtype=torch.float32, device=self.device)
+        old_values = torch.as_tensor(batch["values"], dtype=torch.float32, device=self.device)
+
+        # Standard advantage normalization
+        adv_std = advantages.std()
+        if adv_std > 1e-8:
+            advantages = (advantages - advantages.mean()) / (adv_std + 1e-8)
+
+        # Optional learning rate schedule annealing
+        if ANNEAL_LR:
+            frac = 1.0 - (self.global_steps / max(1, total_steps_target))
+            lr_now = max(1e-6, frac * LEARNING_RATE)
+            for param_group in self.optimizer.param_groups:
+                param_group["lr"] = lr_now
+
+        n = obs.shape[0]
+        minibatch = min(MINIBATCH_SIZE, n)
+        indices = np.arange(n)
+        metrics: Dict[str, List[float]] = {
+            "policy_loss": [],
+            "value_loss": [],
+            "entropy": [],
+            "approx_kl": [],
+            "clipfrac": [],
+            "explained_var": [],
+        }
+
+        for _ in range(PPO_EPOCHS):
+            np.random.shuffle(indices)
+            for start in range(0, n, minibatch):
+                mb = indices[start:start + minibatch]
+                _, new_logprob, entropy, new_value, _ = self.policy.get_action_and_value(obs[mb], actions[mb])
+
+                logratio = new_logprob - old_logprobs[mb]
+                ratio = logratio.exp()
+                mb_adv = advantages[mb]
+
+                # Policy loss (clipped surrogate)
+                pg_loss1 = -mb_adv * ratio
+                pg_loss2 = -mb_adv * torch.clamp(ratio, 1.0 - CLIP_COEF, 1.0 + CLIP_COEF)
+                policy_loss = torch.max(pg_loss1, pg_loss2).mean()
+
+                # Value loss (smooth MSE)
+                value_loss = 0.5 * ((new_value - returns[mb]) ** 2).mean()
+                entropy_loss = entropy.mean()
+
+                loss = policy_loss + VALUE_COEF * value_loss - ENTROPY_COEF * entropy_loss
+
+                self.optimizer.zero_grad(set_to_none=True)
+                loss.backward()
+                nn.utils.clip_grad_norm_(self.policy.parameters(), MAX_GRAD_NORM)
+                self.optimizer.step()
+
+                approx_kl = ((ratio - 1.0) - logratio).mean().detach().cpu().item()
+                clipfrac = ((ratio - 1.0).abs() > CLIP_COEF).float().mean().detach().cpu().item()
+
+                metrics["policy_loss"].append(policy_loss.detach().cpu().item())
+                metrics["value_loss"].append(value_loss.detach().cpu().item())
+                metrics["entropy"].append(entropy_loss.detach().cpu().item())
+                metrics["approx_kl"].append(approx_kl)
+                metrics["clipfrac"].append(clipfrac)
+
+        self.updates += 1
+        y_true = returns.cpu().numpy()
+        y_pred = old_values.cpu().numpy()
+        var_y = np.var(y_true)
+        explained_var = float(1.0 - np.var(y_true - y_pred) / (var_y + 1e-8)) if var_y > 1e-8 else 0.0
+
+        out_metrics = {k: float(np.mean(v)) for k, v in metrics.items() if len(v) > 0}
+        out_metrics["explained_var"] = explained_var
+        return out_metrics
+
+
+# =============================================================================
+# ROLLOUT COLLECTION & GAE
+# =============================================================================
+
+def collect_rollout(
+    policy: ActorCritic,
+    envs: List[PlateBalanceEnv],
+    current_obs: np.ndarray,
+    trainer: PPOTrainer,
+) -> Tuple[Dict[str, np.ndarray], np.ndarray, List[EpisodeStats], int]:
+    """Collect vector rollouts and compute generalized advantage estimations."""
+    n_envs = len(envs)
+    T = ROLLOUT_STEPS
+
+    obs_buf = np.zeros((T, n_envs, OBSERVATION_SIZE), dtype=np.float32)
+    actions_buf = np.zeros((T, n_envs, ACTION_SIZE), dtype=np.float32)
+    raw_actions_buf = np.zeros((T, n_envs, ACTION_SIZE), dtype=np.float32)
+    logprob_buf = np.zeros((T, n_envs), dtype=np.float32)
+    rewards_buf = np.zeros((T, n_envs), dtype=np.float32)
+    terminated_buf = np.zeros((T, n_envs), dtype=np.float32)
+    truncated_buf = np.zeros((T, n_envs), dtype=np.float32)
+    values_buf = np.zeros((T, n_envs), dtype=np.float32)
+
+    completed_episodes: List[EpisodeStats] = []
+    start_steps = trainer.global_steps
+
+    policy.eval()
+    for t in range(T):
+        obs_buf[t] = current_obs
+        obs_t = torch.as_tensor(current_obs, dtype=torch.float32, device=trainer.device)
+        with torch.no_grad():
+            action_t, logprob_t, _, value_t, raw_t = policy.get_action_and_value(obs_t)
+
+        actions = action_t.cpu().numpy()
+        raw_actions = raw_t.cpu().numpy()
+
+        actions_buf[t] = actions
+        raw_actions_buf[t] = raw_actions
+        logprob_buf[t] = logprob_t.cpu().numpy()
+        values_buf[t] = value_t.cpu().numpy()
+
+        next_obs = np.empty_like(current_obs)
+        for e, env in enumerate(envs):
+            next_obs[e], reward, terminated, truncated, info = env.step(actions[e])
+            rewards_buf[t, e] = reward
+            terminated_buf[t, e] = float(terminated)
+            truncated_buf[t, e] = float(truncated)
+            trainer.global_steps += 1
+
+            if terminated or truncated:
+                trainer.episodes += 1
+                trainer.object_counts[info["object"]] = trainer.object_counts.get(info["object"], 0) + 1
+                completed_episodes.append(EpisodeStats(
+                    reward=float(info["episode_reward"]),
+                    length=int(info["episode_length"]),
+                    object_name=str(info["object"]),
+                    success=bool(not info["fallen"] and info["episode_length"] >= MAX_EPISODE_STEPS),
+                    crumbs_on=int(info.get("crumbs_on_plate", 1)),
+                ))
+                # Reset environment immediately upon termination or truncation
+                next_obs[e] = env.reset()
+
+        current_obs = next_obs
+        if trainer.episodes >= NUM_EPISODES:
+            break
+
+    actual_T = t + 1
+    obs_buf = obs_buf[:actual_T]
+    actions_buf = actions_buf[:actual_T]
+    raw_actions_buf = raw_actions_buf[:actual_T]
+    logprob_buf = logprob_buf[:actual_T]
+    rewards_buf = rewards_buf[:actual_T]
+    terminated_buf = terminated_buf[:actual_T]
+    truncated_buf = truncated_buf[:actual_T]
+    values_buf = values_buf[:actual_T]
+
+    with torch.no_grad():
+        next_obs_t = torch.as_tensor(current_obs, dtype=torch.float32, device=trainer.device)
+        next_value = policy.get_value(next_obs_t).cpu().numpy()
+
+    # GAE Computation with proper termination vs. truncation bootstrapping
+    advantages = np.zeros_like(rewards_buf)
+    lastgaelam = np.zeros(n_envs, dtype=np.float32)
+    for t2 in reversed(range(actual_T)):
+        if t2 == actual_T - 1:
+            next_vals = next_value
+        else:
+            next_vals = values_buf[t2 + 1]
+
+        # Only true terminations (falling) zero out future value; timeouts bootstrap value
+        nonterminal = 1.0 - terminated_buf[t2]
+        delta = rewards_buf[t2] + GAMMA * next_vals * nonterminal - values_buf[t2]
+        lastgaelam = delta + GAMMA * GAE_LAMBDA * nonterminal * lastgaelam
+        advantages[t2] = lastgaelam
+
+    returns = advantages + values_buf
+
+    batch = {
+        "obs": obs_buf.reshape(-1, OBSERVATION_SIZE),
+        "actions": actions_buf.reshape(-1, ACTION_SIZE),
+        "raw_actions": raw_actions_buf.reshape(-1, ACTION_SIZE),
+        "logprobs": logprob_buf.reshape(-1),
+        "rewards": rewards_buf.reshape(-1),
+        "terminated": terminated_buf.reshape(-1),
+        "values": values_buf.reshape(-1),
+        "advantages": advantages.reshape(-1),
+        "returns": returns.reshape(-1),
+    }
+    policy.train()
+    return batch, current_obs, completed_episodes, start_steps
+
+
+# =============================================================================
+# EVALUATION & VIDEO VISUALIZATION
+# =============================================================================
+
+def record_video(
+    policy: ActorCritic,
+    out_path: Path,
+    seed: int = 1234,
+    max_steps: int = VIDEO_LENGTH_STEPS,
+    camera: str = "track",
+    specific_object: Optional[str] = None,
+) -> None:
+    """Record and save an MP4 demonstration video of the policy."""
+    if imageio is None:
+        print("[video] imageio / imageio-ffmpeg not installed. Skipping video recording.")
+        return
+
+    env = PlateBalanceEnv(seed=seed, render=True)
+    obs = env.reset(specific_object=specific_object)
+    frames: List[np.ndarray] = []
+    policy.eval()
+
+    try:
+        with torch.no_grad():
+            for _ in range(max_steps):
+                frames.append(env.render(camera=camera))
+                obs_t = torch.as_tensor(obs, dtype=torch.float32, device=DEVICE).unsqueeze(0)
+                action, _, _, _, _ = policy.get_action_and_value(obs_t, deterministic=True)
+                obs, _, terminated, truncated, _ = env.step(action[0].cpu().numpy())
+                if terminated or truncated:
+                    obs = env.reset()
+
+        out_path.parent.mkdir(parents=True, exist_ok=True)
+        imageio.mimsave(out_path, frames, fps=VIDEO_FPS, codec="libx264", quality=7)
+        print(f"[video] Saved {len(frames)} frames to: {out_path}")
+    finally:
+        env.close()
+        policy.train()
+
+
+def evaluate_policy(
+    policy: ActorCritic,
+    episodes_per_object: int = 5,
+    seed: int = 777,
+    deterministic: bool = True,
+    eval_objects: Optional[List[str]] = None,
+) -> Dict[str, Any]:
+    """Benchmark the policy across all 21 supported object geometries."""
+    results: Dict[str, Dict[str, float]] = {}
+    policy.eval()
+
+    active_eval_list = eval_objects or OBJECT_TYPES
+
+    print("\n" + "=" * 88)
+    print("EVALUATION BENCHMARK ACROSS ALL 21 OBJECT GEOMETRIES (INCL. CRUMBLING COOKIE)")
+    print("=" * 88)
+    print(f"{'Category':<14} | {'Object':<16} | {'Reward Mean':<12} | {'Len Mean':<10} | {'Survival %':<12} | {'Tracking Err':<12}")
+    print("-" * 88)
+
+    for obj_name in active_eval_list:
+        category = "Multi-Body" if obj_name == "cookie" else ("Harder" if obj_name in HARDER_OBJECT_TYPES else "Standard")
+        env = PlateBalanceEnv(seed=seed, render=False, active_objects=[obj_name])
+        rewards: List[float] = []
+        lengths: List[int] = []
+        distances: List[float] = []
+        survived = 0
+
+        for ep in range(episodes_per_object):
+            obs = env.reset(specific_object=obj_name)
+            done = False
+            ep_dist = []
+            while not done:
+                with torch.no_grad():
+                    obs_t = torch.as_tensor(obs, dtype=torch.float32, device=DEVICE).unsqueeze(0)
+                    action, _, _, _, _ = policy.get_action_and_value(obs_t, deterministic=deterministic)
+                obs, reward, terminated, truncated, info = env.step(action[0].cpu().numpy())
+                ep_dist.append(info["distance"])
+                if terminated or truncated:
+                    rewards.append(info["episode_reward"])
+                    lengths.append(info["episode_length"])
+                    if not info["fallen"]:
+                        survived += 1
+                    done = True
+
+            distances.append(float(np.mean(ep_dist)))
+
+        env.close()
+        mean_r = float(np.mean(rewards)) if rewards else 0.0
+        mean_l = float(np.mean(lengths)) if lengths else 0.0
+        surv_rate = (survived / episodes_per_object) * 100.0
+        mean_d = float(np.mean(distances)) if distances else 0.0
+
+        results[obj_name] = {
+            "mean_reward": mean_r,
+            "mean_length": mean_l,
+            "survival_rate": surv_rate,
+            "tracking_error": mean_d,
+        }
+        print(f"{category:<14} | {obj_name:<16} | {mean_r:>12.2f} | {mean_l:>10.1f} | {surv_rate:>11.1f}% | {mean_d:>12.4f}m")
+
+    print("=" * 88 + "\n")
+    policy.train()
+    return results
+
+
+def human_view(policy: ActorCritic, specific_object: Optional[str] = None) -> None:
+    """Interactive real-time 3D MuJoCo viewer."""
+    try:
+        import mujoco.viewer
+    except ImportError:
+        print("[viewer] mujoco.viewer is unavailable on this system.")
+        return
+
+    env = PlateBalanceEnv(seed=SEED + 999, render=False)
+    obs = env.reset(specific_object=specific_object)
+    control_dt = PHYSICS_TIMESTEP * CONTROL_DECIMATION
+
+    print(f"\n[viewer] Launching interactive 3D viewer for: {specific_object or 'Random Objects'}")
+    print("[viewer] Controls: Space to pause/resume, Esc/close window to exit.")
+    with mujoco.viewer.launch_passive(env.model, env.data) as viewer:
+        policy.eval()
+        while viewer.is_running():
+            step_start = time.time()
+            with torch.no_grad():
+                obs_t = torch.as_tensor(obs, dtype=torch.float32, device=DEVICE).unsqueeze(0)
+                action, _, _, _, _ = policy.get_action_and_value(obs_t, deterministic=True)
+            obs, _, terminated, truncated, info = env.step(action[0].cpu().numpy())
+            viewer.sync()
+
+            if terminated or truncated:
+                status = "TIMEOUT" if truncated else "FALL"
+                crumbs_info = f" | Crumbs on plate: {info.get('crumbs_on_plate', 1)}/5" if info["object"] == "cookie" else ""
+                print(f"[viewer] Episode end: {status} | Object: {info['object']} | Length: {info['episode_length']}{crumbs_info}")
+                obs = env.reset(specific_object=specific_object)
+
+            # Precise real-time rate pacing
+            elapsed = time.time() - step_start
+            if elapsed < control_dt:
+                time.sleep(control_dt - elapsed)
+
+        policy.train()
+    env.close()
+
+
+def export_onnx(policy: ActorCritic, out_path: Path) -> None:
+    """Export the trained policy backbone and actor head to standard ONNX format."""
+    policy.eval()
+    dummy_input = torch.zeros(1, OBSERVATION_SIZE, dtype=torch.float32, device=DEVICE)
+
+    class ExportWrapper(nn.Module):
+        def __init__(self, p: ActorCritic):
+            super().__init__()
+            self.policy = p
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            mean, _, _ = self.policy(x)
+            return torch.tanh(mean)
+
+    wrapper = ExportWrapper(policy)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    torch.onnx.export(
+        wrapper,
+        dummy_input,
+        str(out_path),
+        input_names=["observation"],
+        output_names=["action"],
+        dynamic_axes={"observation": {0: "batch_size"}, "action": {0: "batch_size"}},
+        opset_version=14,
+    )
+    print(f"[export] Successfully exported ONNX model to: {out_path}")
+    policy.train()
+
+
+# =============================================================================
+# MAIN ENTRY POINT & TRAINING LOOP
+# =============================================================================
+
+def print_config(active_objects: List[str]) -> None:
+    if not PRINT_HYPERPARAMS:
+        return
+    print("=" * 88)
+    print("PlateBalance RL - Advanced MuJoCo + PyTorch PPO System (21 Objects & Cookie Crumble)")
+    print("=" * 88)
+    print(f"Device: {DEVICE} | Master Seed: {SEED}")
+    print(f"Policy: Obs({OBSERVATION_SIZE}) -> {HIDDEN_SIZE} -> {INTERMEDIATE_SIZE} -> {BOTTLENECK_SIZE} -> Act({ACTION_SIZE})")
+    print(f"PPO: Envs={NUM_ENVS} | Rollout={ROLLOUT_STEPS} | Epochs={PPO_EPOCHS} | Batch={MINIBATCH_SIZE} | LR={LEARNING_RATE}")
+    print(f"Target: Episodes={NUM_EPISODES:,} | MaxSteps/Ep={MAX_EPISODE_STEPS:,}")
+    print(f"Active Objects ({len(active_objects)}): {', '.join(active_objects)}")
+    print(f"Domain Randomization: Mass={RANDOMIZE_MASS}, Friction={RANDOMIZE_FRICTION}, Geometry={RANDOMIZE_OBJECT_SIZE}")
+    print(f"Actuator Torque: {MAX_PLATE_TORQUE} N*m | Control Rate: {1.0/(PHYSICS_TIMESTEP*CONTROL_DECIMATION):.0f} Hz")
+    print("=" * 88)
+
+
+def main() -> None:
+    global SEED, DEVICE, NUM_EPISODES
+    parser = argparse.ArgumentParser(description="PlateBalance RL: MuJoCo PPO Object Balancing")
+    parser.add_argument("--eval", action="store_true", help="Run benchmark evaluation across all 21 object shapes")
+    parser.add_argument("--human-view", action="store_true", help="Launch interactive 3D viewer")
+    parser.add_argument("--record-video", action="store_true", help="Record evaluation video")
+    parser.add_argument("--export-onnx", type=str, default="", help="Export model to ONNX file path")
+    parser.add_argument("--checkpoint", type=str, default="", help="Path to checkpoint folder to load")
+    parser.add_argument("--resume", action="store_true", help="Auto-resume training from latest checkpoint")
+    parser.add_argument("--episodes", type=int, default=NUM_EPISODES, help="Total training episodes")
+    parser.add_argument("--seed", type=int, default=SEED, help="Random seed")
+    parser.add_argument("--device", type=str, default=DEVICE, help="Compute device (cuda or cpu)")
+    parser.add_argument("--object", type=str, default="", help="Specific object to balance for training/eval/viewer")
+    parser.add_argument("--category", type=str, default="all", choices=["all", "standard", "harder", "cookie"],
+                        help="Filter object category (all, standard, harder, cookie)")
+    args = parser.parse_args()
+
+    SEED = args.seed
+    DEVICE = args.device
+    NUM_EPISODES = args.episodes
+
+    # Filter active objects
+    if args.object:
+        if args.object not in OBJECT_INFO:
+            raise ValueError(f"Unknown object '{args.object}'. Available: {', '.join(OBJECT_TYPES)}")
+        active_objects = [args.object]
+    elif args.category == "standard":
+        active_objects = STANDARD_OBJECT_TYPES
+    elif args.category == "harder":
+        active_objects = HARDER_OBJECT_TYPES
+    elif args.category == "cookie":
+        active_objects = MULTI_BODY_OBJECT_TYPES
+    else:
+        active_objects = OBJECT_TYPES
+
+    set_seed(SEED)
+    print_config(active_objects)
+
+    root = Path(OUTPUT_DIR) / EXPERIMENT_NAME
+    root.mkdir(parents=True, exist_ok=True)
+
+    trainer = PPOTrainer(device=DEVICE)
+
+    # Determine checkpoint loading
+    ckpt_path: Optional[Path] = None
+    if args.checkpoint:
+        ckpt_path = Path(args.checkpoint)
+    elif args.resume or RESUME:
+        ckpt_path = Path(RESUME_PATH) if RESUME_PATH else latest_checkpoint(root / "checkpoints")
+
+    if ckpt_path is not None and ckpt_path.exists():
+        trainer.load(ckpt_path)
+
+    # Export ONNX mode
+    if args.export_onnx:
+        export_onnx(trainer.policy, Path(args.export_onnx))
+        return
+
+    # Interactive Human Viewer mode
+    if args.human_view or HUMAN_VIEW:
+        human_view(trainer.policy, specific_object=args.object or None)
+        return
+
+    # Benchmark Evaluation mode
+    if args.eval:
+        evaluate_policy(trainer.policy, eval_objects=active_objects)
+        return
+
+    # Record Video mode
+    if args.record_video:
+        vid_path = root / "eval_videos" / "eval_demo.mp4"
+        record_video(trainer.policy, vid_path, seed=SEED, specific_object=args.object or None)
+        return
+
+    # Standard Training Setup
+    envs = [
+        PlateBalanceEnv(SEED + 1000 * i + trainer.global_steps, render=False, active_objects=active_objects)
+        for i in range(NUM_ENVS)
+    ]
+    current_obs = np.stack([env.reset() for env in envs], axis=0)
+
+    start_time = time.time()
+    running_rewards: List[float] = []
+    running_lengths: List[int] = []
+    running_success: List[float] = []
+
+    print(f"\n[train] Starting PPO training loop with {NUM_ENVS} parallel environments...\n")
+
+    try:
+        while trainer.episodes < NUM_EPISODES:
+            batch, current_obs, completed, rollout_start = collect_rollout(
+                trainer.policy, envs, current_obs, trainer
+            )
+
+            if batch["obs"].shape[0] >= 2:
+                metrics = trainer.update(batch, total_steps_target=NUM_EPISODES * 500)
+            else:
+                metrics = {
+                    "policy_loss": 0.0, "value_loss": 0.0, "entropy": 0.0,
+                    "approx_kl": 0.0, "clipfrac": 0.0, "explained_var": 0.0,
+                }
+
+            running_rewards.extend(ep.reward for ep in completed)
+            running_lengths.extend(ep.length for ep in completed)
+            running_success.extend(1.0 if ep.success else 0.0 for ep in completed)
+
+            # Periodic console logging
+            if trainer.global_steps - trainer.last_log_step >= LOGGING_STEPS:
+                trainer.last_log_step = trainer.global_steps
+                elapsed = max(time.time() - start_time, 1e-9)
+                sps = (trainer.global_steps / elapsed)
+                mean_reward = float(np.mean(running_rewards[-50:])) if running_rewards else 0.0
+                mean_len = float(np.mean(running_lengths[-50:])) if running_lengths else 0.0
+                succ_rate = (float(np.mean(running_success[-50:])) * 100.0) if running_success else 0.0
+
+                print(
+                    f"step={trainer.global_steps:>9,} | ep={trainer.episodes:>7,} | "
+                    f"sps={sps:>6.0f} | r50={mean_reward:>8.2f} | len50={mean_len:>6.1f} | "
+                    f"succ={succ_rate:>5.1f}% | pi={metrics.get('policy_loss', 0.0):+.4f} | "
+                    f"vf={metrics.get('value_loss', 0.0):.4f} | kl={metrics.get('approx_kl', 0.0):.5f}"
+                )
+
+            # Periodic checkpoint saving
+            if trainer.global_steps - trainer.last_save_step >= SAVE_STEPS:
+                trainer.last_save_step = trainer.global_steps
+                ckpt = root / "checkpoints" / f"step_{trainer.global_steps:08d}"
+                trainer.save(ckpt, latest=True, extra={
+                    "mean_reward_50": float(np.mean(running_rewards[-50:])) if running_rewards else 0.0,
+                    "mean_length_50": float(np.mean(running_lengths[-50:])) if running_lengths else 0.0,
+                })
+                print(f"[save] Checkpoint saved: {ckpt}")
+
+            # Periodic visualization video recording
+            if VISUALIZE and trainer.global_steps - trainer.last_video_step >= VIDEO_EVERY_STEPS:
+                trainer.last_video_step = trainer.global_steps
+                ckpt = root / "checkpoints" / f"step_{trainer.global_steps:08d}"
+                ckpt.mkdir(parents=True, exist_ok=True)
+                video_path = ckpt / f"balance_step_{trainer.global_steps:08d}.mp4"
+                print(f"[video] Recording rollout video to: {video_path}")
+                record_video(trainer.policy, video_path, seed=SEED + trainer.global_steps, max_steps=VIDEO_LENGTH_STEPS)
+                trainer.save(ckpt, latest=True, extra={"video": str(video_path.name)})
+
+    except KeyboardInterrupt:
+        print("\n[interrupt] Training interrupted by user. Saving emergency checkpoint...")
+        ckpt = root / "checkpoints" / f"step_{trainer.global_steps:08d}_interrupt"
+        trainer.save(ckpt, latest=True, extra={"interrupted": True})
+        print(f"[interrupt] Saved emergency checkpoint: {ckpt}")
+    finally:
+        for env in envs:
+            env.close()
+
+    final_dir = root / "checkpoints" / f"step_{trainer.global_steps:08d}_final"
+    trainer.save(final_dir, latest=True, extra={"finished": True})
+    print(f"\n[done] Training completed! Total episodes={trainer.episodes:,} steps={trainer.global_steps:,}")
+    print(f"[done] Final model and state saved to: {final_dir}")
+
+
+if __name__ == "__main__":
+    main()

inference.py

@@ -0,0 +1,474 @@
+#!/usr/bin/env python3
+"""
+Test MrBalance directly from Hugging Face Hub.
+
+Requirements:
+    pip install -U torch transformers safetensors mujoco
+
+Files:
+    balance_plate_rl.py
+    test_mrbalance_hf.py
+
+Usage:
+    python test_mrbalance_hf.py
+    python test_mrbalance_hf.py --object sphere
+    python test_mrbalance_hf.py --object egg
+    python test_mrbalance_hf.py --object heavy_ball
+
+The script:
+    1. Downloads MrBalance from Hugging Face.
+    2. Loads it through AutoModel with trust_remote_code=True.
+    3. Creates the original MuJoCo environment.
+    4. Uses the Hugging Face policy to control the plate.
+    5. Prints episode statistics.
+    6. Optionally launches the MuJoCo viewer.
+"""
+
+from __future__ import annotations
+
+import argparse
+import time
+
+import numpy as np
+import torch
+from transformers import AutoModel
+
+from balance_plate_rl import (
+    DEVICE,
+    MAX_EPISODE_STEPS,
+    OBJECT_TYPES,
+    PlateBalanceEnv,
+)
+
+
+# ============================================================================
+# CONFIG
+# ============================================================================
+
+MODEL_ID = "fromziro/MrBalance"
+
+DEFAULT_EPISODES = 5
+DEFAULT_OBJECT = "sphere"
+
+DEVICE_TO_USE = torch.device(
+    DEVICE
+)
+
+
+# ============================================================================
+# MODEL LOADING
+# ============================================================================
+
+def load_model():
+    print("=" * 72)
+    print("Loading MrBalance from Hugging Face")
+    print("=" * 72)
+
+    print(f"Model:  {MODEL_ID}")
+    print(f"Device: {DEVICE_TO_USE}")
+    print()
+
+    model = AutoModel.from_pretrained(
+        MODEL_ID,
+        trust_remote_code=True,
+    )
+
+    model = model.to(DEVICE_TO_USE)
+    model.eval()
+
+    print("[load] Model loaded successfully.")
+    print()
+
+    # Print basic architecture information.
+    print(
+        f"[model] Observation size: "
+        f"{model.config.observation_size}"
+    )
+
+    print(
+        f"[model] Hidden size:       "
+        f"{model.config.hidden_size}"
+    )
+
+    print(
+        f"[model] Intermediate size: "
+        f"{model.config.intermediate_size}"
+    )
+
+    print(
+        f"[model] Bottleneck size:   "
+        f"{model.config.bottleneck_size}"
+    )
+
+    print(
+        f"[model] Action size:       "
+        f"{model.config.action_size}"
+    )
+
+    print()
+
+    return model
+
+
+# ============================================================================
+# SINGLE EPISODE
+# ============================================================================
+
+def run_episode(
+    model,
+    object_name: str,
+    seed: int,
+    render: bool = False,
+):
+    env = PlateBalanceEnv(
+        seed=seed,
+        render=render,
+        active_objects=[object_name],
+    )
+
+    obs = env.reset(
+        specific_object=object_name
+    )
+
+    total_reward = 0.0
+    episode_length = 0
+
+    if render:
+        try:
+            import mujoco.viewer
+
+            viewer_context = (
+                mujoco.viewer.launch_passive(
+                    env.model,
+                    env.data,
+                )
+            )
+        except Exception as exc:
+            env.close()
+            raise RuntimeError(
+                f"Could not launch MuJoCo viewer: {exc}"
+            ) from exc
+    else:
+        viewer_context = None
+
+    try:
+        if viewer_context is not None:
+            with viewer_context as viewer:
+
+                while viewer.is_running():
+
+                    step_start = time.time()
+
+                    obs_tensor = torch.as_tensor(
+                        obs,
+                        dtype=torch.float32,
+                        device=DEVICE_TO_USE,
+                    ).unsqueeze(0)
+
+                    with torch.no_grad():
+                        output = model(
+                            obs_tensor,
+                            deterministic=True,
+                        )
+
+                    action = (
+                        output.action[0]
+                        .detach()
+                        .cpu()
+                        .numpy()
+                        .astype(np.float64)
+                    )
+
+                    obs, reward, terminated, truncated, info = (
+                        env.step(action)
+                    )
+
+                    total_reward += reward
+                    episode_length += 1
+
+                    viewer.sync()
+
+                    if terminated or truncated:
+                        break
+
+                    # Match roughly the environment's 100 Hz control rate.
+                    target_dt = 0.01
+                    elapsed = time.time() - step_start
+
+                    if elapsed < target_dt:
+                        time.sleep(
+                            target_dt - elapsed
+                        )
+
+        else:
+            while True:
+
+                obs_tensor = torch.as_tensor(
+                    obs,
+                    dtype=torch.float32,
+                    device=DEVICE_TO_USE,
+                ).unsqueeze(0)
+
+                with torch.no_grad():
+                    output = model(
+                        obs_tensor,
+                        deterministic=True,
+                    )
+
+                action = (
+                    output.action[0]
+                    .detach()
+                    .cpu()
+                    .numpy()
+                    .astype(np.float64)
+                )
+
+                obs, reward, terminated, truncated, info = (
+                    env.step(action)
+                )
+
+                total_reward += reward
+                episode_length += 1
+
+                if terminated or truncated:
+                    break
+
+    finally:
+        env.close()
+
+    success = (
+        not info["fallen"]
+        and episode_length >= MAX_EPISODE_STEPS
+    )
+
+    return {
+        "object": object_name,
+        "reward": float(total_reward),
+        "length": int(episode_length),
+        "success": bool(success),
+        "fallen": bool(info["fallen"]),
+        "distance": float(info["distance"]),
+        "velocity": float(info["object_velocity"]),
+    }
+
+
+# ============================================================================
+# RANDOM OBSERVATION SANITY CHECK
+# ============================================================================
+
+def sanity_check(model):
+    """
+    Verify that the Hugging Face model can actually execute inference
+    independently of MuJoCo.
+    """
+
+    print("=" * 72)
+    print("Running model sanity check")
+    print("=" * 72)
+
+    x = torch.randn(
+        4,
+        64,
+        dtype=torch.float32,
+        device=DEVICE_TO_USE,
+    )
+
+    with torch.no_grad():
+        output = model(
+            x,
+            deterministic=True,
+        )
+
+    print(
+        "[sanity] action shape:",
+        tuple(output.action.shape),
+    )
+
+    print(
+        "[sanity] value shape:",
+        tuple(output.value.shape),
+    )
+
+    print(
+        "[sanity] action range:",
+        float(output.action.min()),
+        "to",
+        float(output.action.max()),
+    )
+
+    print(
+        "[sanity] example action:",
+        output.action[0].detach().cpu().numpy(),
+    )
+
+    print("[sanity] PASS")
+    print()
+
+
+# ============================================================================
+# MAIN
+# ============================================================================
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Test MrBalance from Hugging Face."
+    )
+
+    parser.add_argument(
+        "--object",
+        type=str,
+        default=DEFAULT_OBJECT,
+        choices=OBJECT_TYPES,
+        help="Object to balance.",
+    )
+
+    parser.add_argument(
+        "--episodes",
+        type=int,
+        default=DEFAULT_EPISODES,
+        help="Number of evaluation episodes.",
+    )
+
+    parser.add_argument(
+        "--render",
+        action="store_true",
+        help="Launch interactive MuJoCo viewer.",
+    )
+
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=12345,
+        help="Evaluation seed.",
+    )
+
+    args = parser.parse_args()
+
+    model = load_model()
+
+    sanity_check(model)
+
+    print("=" * 72)
+    print(
+        f"Testing object: {args.object}"
+    )
+    print(
+        f"Episodes:       {args.episodes}"
+    )
+    print("=" * 72)
+    print()
+
+    results = []
+
+    for episode in range(args.episodes):
+
+        print(
+            f"[episode {episode + 1}/{args.episodes}] "
+            f"Running..."
+        )
+
+        result = run_episode(
+            model=model,
+            object_name=args.object,
+            seed=args.seed + episode,
+            render=args.render,
+        )
+
+        results.append(result)
+
+        status = (
+            "SUCCESS"
+            if result["success"]
+            else (
+                "FALL"
+                if result["fallen"]
+                else "TIMEOUT"
+            )
+        )
+
+        print(
+            f"  status:   {status}"
+        )
+
+        print(
+            f"  reward:   {result['reward']:.2f}"
+        )
+
+        print(
+            f"  length:   {result['length']}"
+        )
+
+        print(
+            f"  distance: {result['distance']:.4f} m"
+        )
+
+        print(
+            f"  velocity: {result['velocity']:.4f} m/s"
+        )
+
+        print()
+
+    # ------------------------------------------------------------------------
+    # Summary
+    # ------------------------------------------------------------------------
+
+    rewards = [
+        r["reward"]
+        for r in results
+    ]
+
+    lengths = [
+        r["length"]
+        for r in results
+    ]
+
+    distances = [
+        r["distance"]
+        for r in results
+    ]
+
+    successes = sum(
+        r["success"]
+        for r in results
+    )
+
+    print("=" * 72)
+    print("RESULTS")
+    print("=" * 72)
+
+    print(
+        f"Object:            {args.object}"
+    )
+
+    print(
+        f"Reward mean:       {np.mean(rewards):.2f}"
+    )
+
+    print(
+        f"Reward std:        {np.std(rewards):.2f}"
+    )
+
+    print(
+        f"Length mean:       {np.mean(lengths):.1f}"
+    )
+
+    print(
+        f"Survival rate:     "
+        f"{100.0 * sum(not r['fallen'] for r in results) / len(results):.1f}%"
+    )
+
+    print(
+        f"Success rate:      "
+        f"{100.0 * successes / len(results):.1f}%"
+    )
+
+    print(
+        f"Tracking error:    "
+        f"{np.mean(distances):.4f} m"
+    )
+
+    print("=" * 72)
+
+
+if __name__ == "__main__":
+    main()