Spaces:
Running
Running
| """Mr. Balance — interactive MuJoCo demo of the fromziro/MrBalance RL policy. | |
| Mr. Balance is a 2-axis gimbal plate driven by a 33k-parameter PPO actor-critic. | |
| This Space rolls out the released policy in the authors' own MuJoCo environment | |
| (`balance_plate_rl.py`, vendored unmodified from the model repo) and renders the | |
| episode to an mp4. | |
| """ | |
| from __future__ import annotations | |
| import math | |
| import os | |
| import random | |
| import sys | |
| import tempfile | |
| import time | |
| import numpy as np | |
| import torch | |
| from transformers import AutoModel | |
| sys.stdout.reconfigure(line_buffering=True) | |
| sys.stderr.reconfigure(line_buffering=True) | |
| MODEL_ID = "fromziro/MrBalance" | |
| # The policy is a 33k-parameter MLP: CPU is the right device, and keeping the | |
| # thread count low avoids oversubscribing the shared CPU while rendering. | |
| torch.set_num_threads(2) | |
| policy = AutoModel.from_pretrained(MODEL_ID, trust_remote_code=True) | |
| policy.eval() | |
| print( | |
| f"[boot] loaded {MODEL_ID}: " | |
| f"{sum(p.numel() for p in policy.parameters()):,} params, " | |
| f"obs={policy.config.observation_size}, act={policy.config.action_size}" | |
| ) | |
| # IMPORTANT: torch and transformers must be imported (and the policy loaded) | |
| # BEFORE mujoco. Importing mujoco first pulls in OSMesa's software GL stack, | |
| # whose bundled LLVM clashes with the one inside torch and hangs / segfaults | |
| # the process (verified on this Space: exit code 139 with the reverse order). | |
| # | |
| # MuJoCo needs a headless GL backend. OSMesa is pure software rendering, so it | |
| # works on any hardware tier (this demo is CPU-only: all the work is physics + | |
| # rasterization). Override with the MUJOCO_GL space variable if a faster | |
| # backend (e.g. egl) is available. | |
| os.environ.setdefault("MUJOCO_GL", "osmesa") | |
| def _cpu_quota() -> int: | |
| """Number of CPUs this container may actually use (cgroup v2 quota).""" | |
| try: | |
| quota, period = open("/sys/fs/cgroup/cpu.max").read().split() | |
| if quota != "max": | |
| return max(1, int(int(quota) / int(period))) | |
| except Exception: | |
| pass | |
| return max(1, len(os.sched_getaffinity(0))) | |
| CPUS = _cpu_quota() | |
| # llvmpipe (OSMesa's software rasterizer) otherwise spawns one thread per | |
| # *host* core and thrashes against the container's much smaller CPU quota. | |
| os.environ.setdefault("LP_NUM_THREADS", str(CPUS)) | |
| print(f"[boot] cpu quota={CPUS}, LP_NUM_THREADS={os.environ['LP_NUM_THREADS']}") | |
| import gradio as gr # noqa: E402 | |
| import imageio.v2 as imageio # noqa: E402 | |
| import mujoco # noqa: E402 | |
| from balance_plate_rl import ( # noqa: E402 | |
| CONTROL_DECIMATION, | |
| MAX_PLATE_ANGLE, | |
| OBJECT_INFO, | |
| PHYSICS_TIMESTEP, | |
| PlateBalanceEnv, | |
| ) | |
| # 100 Hz RL control (400 Hz physics / 4 substeps). | |
| CONTROL_HZ = round(1.0 / (PHYSICS_TIMESTEP * CONTROL_DECIMATION)) | |
| RENDER_EVERY = 4 # -> 25 rendered frames per simulated second | |
| VIDEO_FPS = CONTROL_HZ // RENDER_EVERY # real-time playback | |
| POST_FALL_STEPS = 100 # keep filming ~1 s after a fall | |
| TRAINED_ON = ("sphere", "egg", "heavy_ball") | |
| # Ordered for the dropdown: the three training objects first, then the | |
| # zero-shot generalization objects from the model card's table. | |
| OBJECTS = [ | |
| "sphere", "egg", "heavy_ball", | |
| "disk", "cup", "coin", "stick", "tall", "triangle", "block", "puck", | |
| "cone", "capsule", "wedge", "tetra", "flat_bar", "cross", "L_shape", | |
| "wide_block", "offcenter_block", "cookie", | |
| ] | |
| PRETTY = { | |
| "sphere": "Sphere", "egg": "Egg", "heavy_ball": "Heavy bowling ball", | |
| "disk": "Disk", "cup": "Hollow cup", "coin": "Coin", "stick": "Standing stick", | |
| "tall": "Tall block", "triangle": "Triangular prism", "block": "Cube", | |
| "puck": "Puck", "cone": "Cone", "capsule": "Lying capsule", "wedge": "Ramp wedge", | |
| "tetra": "Tetrahedron", "flat_bar": "Long flat bar", "cross": "Cross / plus", | |
| "L_shape": "Asymmetric L-shape", "wide_block": "Wide tile", | |
| "offcenter_block": "Off-center mass block", "cookie": "Crumbling cookie (5 crumbs)", | |
| } | |
| def _label(name: str) -> str: | |
| tag = "trained" if name in TRAINED_ON else "zero-shot" | |
| if name == "cookie": | |
| tag = "zero-shot, brutal" | |
| return f"{PRETTY[name]} — {tag}" | |
| OBJECT_CHOICES = [(_label(n), n) for n in OBJECTS] | |
| CAMERAS = ["track", "isometric", "top_down", "side"] | |
| RESOLUTIONS = {"480 x 360": (480, 360), "640 x 480": (640, 480), "320 x 240": (320, 240)} | |
| QUALITIES = ["Fast", "Pretty (anti-aliased + shadows, slower)"] | |
| def _kick(env: PlateBalanceEnv, speed: float, rng: np.random.Generator) -> None: | |
| """Shove the object sideways with a random horizontal impulse (m/s).""" | |
| angle = float(rng.uniform(0.0, 2.0 * math.pi)) | |
| dv = np.array([speed * math.cos(angle), speed * math.sin(angle)]) | |
| if env.current_object == "cookie": | |
| for dof in env.crumb_dof_addrs: | |
| env.data.qvel[dof:dof + 2] += dv | |
| else: | |
| dof = env.object_dof_addr | |
| env.data.qvel[dof:dof + 2] += dv | |
| def simulate( | |
| object_name: str, | |
| nudge: float = 0.0, | |
| seconds: float = 5.0, | |
| camera: str = "track", | |
| resolution: str = "480 x 360", | |
| quality: str = "Fast", | |
| seed: int = 12345, | |
| randomize_seed: bool = True, | |
| ): | |
| """Roll out the MrBalance policy on one object and render the episode to mp4. | |
| Args: | |
| object_name: which object to drop on the plate (e.g. "sphere", "cookie"). | |
| nudge: strength of periodic random sideways shoves, in m/s (0 = none). | |
| seconds: simulated seconds to roll out. | |
| camera: MuJoCo camera to film from ("track", "isometric", "top_down", "side"). | |
| resolution: rendered video size. | |
| quality: "Fast" (no anti-aliasing or shadows) or "Pretty ..." (slower). | |
| seed: RNG seed (controls spawn pose plus the randomized mass/size/friction). | |
| randomize_seed: draw a fresh random seed instead of using `seed`. | |
| Returns: | |
| (mp4 path, markdown episode report, the seed that was actually used). | |
| """ | |
| if object_name not in OBJECT_INFO: | |
| raise gr.Error(f"Unknown object: {object_name}") | |
| if randomize_seed: | |
| seed = random.randint(0, 2**31 - 1) | |
| seed = int(seed) | |
| width, height = RESOLUTIONS.get(resolution, (480, 360)) | |
| seconds = float(np.clip(seconds, 2.0, 15.0)) | |
| control_steps = int(seconds * CONTROL_HZ) | |
| rng = np.random.default_rng(seed) | |
| env = PlateBalanceEnv(seed=seed, render=False, active_objects=[object_name]) | |
| obs = env.reset(specific_object=object_name) | |
| # MuJoCo's offscreen buffer defaults to 640x480; size it to the request. | |
| env.model.vis.global_.offwidth = width | |
| env.model.vis.global_.offheight = height | |
| # Everything is rasterized in software (llvmpipe), so 4x multisampling and | |
| # a 4096px shadow map cost more than the physics does. "Fast" turns them | |
| # off and is ~3x quicker; "Pretty" keeps the authors' default look. | |
| fast = not str(quality).startswith("Pretty") | |
| env.model.vis.quality.offsamples = 0 if fast else 4 | |
| env.model.vis.quality.shadowsize = 0 if fast else 2048 | |
| renderer = mujoco.Renderer(env.model, height=height, width=width) | |
| if fast: | |
| renderer.scene.flags[mujoco.mjtRndFlag.mjRND_SHADOW] = 0 | |
| renderer.scene.flags[mujoco.mjtRndFlag.mjRND_REFLECTION] = 0 | |
| frames: list[np.ndarray] = [] | |
| distances: list[float] = [] | |
| total_reward = 0.0 | |
| steps = 0 | |
| fell_at = None | |
| kicks = 0 | |
| kick_period = 2 * CONTROL_HZ # a shove every 2 simulated seconds | |
| first_kick = int(1.5 * CONTROL_HZ) | |
| extra = 0 | |
| info: dict = {"fallen": False, "distance": 0.0, "object_velocity": 0.0} | |
| t0 = time.perf_counter() | |
| render_time = 0.0 | |
| try: | |
| with torch.no_grad(): | |
| for t in range(control_steps): | |
| if ( | |
| nudge > 0.0 | |
| and fell_at is None | |
| and t >= first_kick | |
| and (t - first_kick) % kick_period == 0 | |
| ): | |
| _kick(env, float(nudge), rng) | |
| kicks += 1 | |
| obs_t = torch.as_tensor(obs, dtype=torch.float32).unsqueeze(0) | |
| action = ( | |
| policy(obs_t, deterministic=True) | |
| .action[0] | |
| .numpy() | |
| .astype(np.float64) | |
| ) | |
| obs, reward, terminated, truncated, info = env.step(action) | |
| if fell_at is None: | |
| # Post-fall frames are only filmed for the visual tumble; | |
| # they must not pollute the episode statistics. | |
| total_reward += reward | |
| distances.append(info["distance"]) | |
| steps += 1 | |
| if t % RENDER_EVERY == 0: | |
| r0 = time.perf_counter() | |
| renderer.update_scene(env.data, camera=camera) | |
| frames.append(renderer.render().copy()) | |
| render_time += time.perf_counter() - r0 | |
| if terminated and fell_at is None: | |
| fell_at = steps | |
| if fell_at is not None: | |
| extra += 1 | |
| if extra >= POST_FALL_STEPS: # film the tumble, then stop | |
| break | |
| if truncated: | |
| break | |
| finally: | |
| renderer.close() | |
| env.close() | |
| path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name | |
| imageio.mimsave( | |
| path, | |
| frames, | |
| fps=VIDEO_FPS, | |
| codec="libx264", | |
| quality=7, | |
| macro_block_size=1, | |
| ) | |
| wall = time.perf_counter() - t0 | |
| tilt = math.degrees( | |
| float( | |
| np.linalg.norm( | |
| [ | |
| env.data.qpos[env.plate_roll_qpos], | |
| env.data.qpos[env.plate_pitch_qpos], | |
| ] | |
| ) | |
| ) | |
| ) | |
| survived = fell_at is None | |
| verdict = ( | |
| f"✅ **Balanced** for the full {steps / CONTROL_HZ:.1f} s" | |
| if survived | |
| else f"❌ **Dropped it** after {fell_at / CONTROL_HZ:.2f} s" | |
| ) | |
| print( | |
| f"[run] {object_name} seed={seed} nudge={nudge} {width}x{height} " | |
| f"q={'fast' if fast else 'pretty'} steps={steps} frames={len(frames)} " | |
| f"survived={survived} wall={wall:.1f}s render={render_time:.1f}s " | |
| f"({1000 * render_time / max(1, len(frames)):.0f} ms/frame)" | |
| ) | |
| report = f"""### {PRETTY[object_name]} — {'trained on' if object_name in TRAINED_ON else 'zero-shot'} | |
| {verdict} | |
| | | | | |
| |---|---| | |
| | Control steps | {steps} @ {CONTROL_HZ} Hz | | |
| | Total reward | {total_reward:,.1f} | | |
| | Mean tracking error | {float(np.mean(distances)):.3f} m | | |
| | Max tracking error | {float(np.max(distances)):.3f} m | | |
| | Final plate tilt | {tilt:.2f}° (limit {math.degrees(MAX_PLATE_ANGLE):.0f}°) | | |
| | Sampled mass | {env.object_mass:.2f} kg | | |
| | Sampled size / friction | ×{env.object_scale:.2f} / {env.base_friction:.2f} | | |
| | Random shoves applied | {kicks}{f' @ {nudge:.1f} m/s' if kicks else ''} | | |
| | Seed | {seed} | | |
| | Wall clock | {wall:.1f} s ({len(frames)} frames, {render_time:.1f} s rendering) | | |
| """ | |
| return path, report, seed | |
| CSS = """ | |
| #col-container { max-width: 1200px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| """ | |
| with gr.Blocks(title="Mr. Balance") as demo: | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown( | |
| """# 🍽️ Mr. Balance | |
| Mr. Balance is a plate. His only goal is to keep whatever you put on him from falling off. | |
| Pick an object, optionally shove it around, and watch the released | |
| [`fromziro/MrBalance`](https://huggingface.co/fromziro/MrBalance) policy — a 33k-parameter | |
| PPO actor–critic outputting roll/pitch torques at 100 Hz — fight to keep it centred in MuJoCo. | |
| It was trained on only **3** objects (sphere, egg, heavy ball); everything else is zero-shot. | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| object_dd = gr.Dropdown( | |
| label="Object on the plate", | |
| choices=OBJECT_CHOICES, | |
| value="sphere", | |
| ) | |
| nudge = gr.Slider( | |
| label="Random shoves (m/s, every 2 s)", | |
| minimum=0.0, | |
| maximum=2.0, | |
| step=0.1, | |
| value=0.0, | |
| info="0 = leave him alone. Higher = keep kicking the object sideways.", | |
| ) | |
| run_btn = gr.Button("Balance it 🍽️", variant="primary") | |
| with gr.Accordion("Advanced settings", open=False): | |
| seconds = gr.Slider( | |
| label="Simulated seconds", | |
| minimum=2.0, | |
| maximum=15.0, | |
| step=1.0, | |
| value=5.0, | |
| ) | |
| camera = gr.Dropdown( | |
| label="Camera", choices=CAMERAS, value="track" | |
| ) | |
| resolution = gr.Dropdown( | |
| label="Video size", | |
| choices=list(RESOLUTIONS), | |
| value="480 x 360", | |
| ) | |
| quality = gr.Dropdown( | |
| label="Render quality", | |
| choices=QUALITIES, | |
| value="Fast", | |
| info="Rendering is done in software on CPU, so " | |
| "anti-aliasing and shadows cost real time.", | |
| ) | |
| randomize_seed = gr.Checkbox(label="Randomize seed", value=True) | |
| seed = gr.Number(label="Seed", value=12345, precision=0) | |
| with gr.Column(scale=2): | |
| video = gr.Video( | |
| label="Episode", | |
| autoplay=True, | |
| loop=True, | |
| height=430, | |
| ) | |
| report = gr.Markdown() | |
| gr.Examples( | |
| examples=[ | |
| ["sphere", 0.0], | |
| ["egg", 0.8], | |
| ["heavy_ball", 1.0], | |
| ["cookie", 0.0], | |
| ["stick", 0.0], | |
| ["cone", 1.0], | |
| ["flat_bar", 0.6], | |
| ["cup", 1.2], | |
| ], | |
| inputs=[object_dd, nudge], | |
| outputs=[video, report, seed], | |
| fn=simulate, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| label="Examples (object, shove strength)", | |
| ) | |
| with gr.Accordion("The authors' original demo video", open=False): | |
| gr.Video( | |
| value="assets/mrbalance_official.mp4", | |
| label="fromziro/MrBalance — assets/video.mp4 (Apache-2.0)", | |
| autoplay=False, | |
| loop=True, | |
| interactive=False, | |
| ) | |
| gr.Markdown( | |
| "Physics, environment and policy by **FromZero** " | |
| "(Paul Courneya, Jonathon LY, User110), Apache-2.0. " | |
| "This Space runs the authors' `balance_plate_rl.py` environment unmodified; " | |
| "the *random shoves* slider is an extra stress test added here." | |
| ) | |
| inputs = [ | |
| object_dd, nudge, seconds, camera, resolution, quality, seed, randomize_seed | |
| ] | |
| run_btn.click( | |
| simulate, inputs=inputs, outputs=[video, report, seed], api_name="simulate" | |
| ) | |
| if __name__ == "__main__": | |
| demo.queue(default_concurrency_limit=1, max_size=24).launch( | |
| theme=gr.themes.Citrus(), css=CSS, mcp_server=True | |
| ) | |