"""Independent explanatory experiments; NOT a reproduction of SuperMimic. CPU MuJoCo physics, synthetic references, no pretrained models or motion datasets. The G1 model is from Unitree's unitree_mujoco checkout. EGL only renders frames. """ import argparse import hashlib import json import math import os from pathlib import Path import platform import subprocess import numpy as np import mujoco def encode(path, frames, fps=24): h, w = frames[0].shape[:2] p = subprocess.Popen(['ffmpeg', '-nostdin', '-y', '-v', 'error', '-f', 'rawvideo', '-pix_fmt', 'rgb24', '-s', f'{w}x{h}', '-r', str(fps), '-i', '-', '-an', '-c:v', 'libx264', '-threads', '2', '-preset', 'fast', '-crf', '19', '-pix_fmt', 'yuv420p', '-movflags', '+faststart', str(path)], stdin=subprocess.PIPE) for frame in frames: p.stdin.write(frame.tobytes()) p.stdin.close() if p.wait(): raise RuntimeError('Video encoding failed') def g1(out, model_path): model = mujoco.MjModel.from_xml_path(str(model_path)) model.opt.timestep = .002 model.vis.global_.offwidth = 960 model.vis.global_.offheight = 540 renderer = mujoco.Renderer(model, height=540, width=960) names = [mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_ACTUATOR, i) for i in range(model.nu)] joints = [model.actuator_trnid[i, 0] for i in range(model.nu)] qids = np.array([model.jnt_qposadr[j] for j in joints]) vids = np.array([model.jnt_dofadr[j] for j in joints]) q0 = np.zeros(model.nu) for i, j in enumerate(joints): name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_JOINT, j) if 'hip_pitch' in name: q0[i] = -.1 elif 'knee' in name: q0[i] = .2 elif 'ankle_pitch' in name: q0[i] = -.1 elif 'elbow' in name: q0[i] = .6 elif 'shoulder_roll' in name: q0[i] = .12 if 'left' in name else -.12 cases = [('reference', 1., .35, 0., False), ('pd_slow', 1., .35, 0., True), ('pd_fast', 2.5, .8, 0., True), ('pd_push', 1., .35, 90., True), ('pd_soft', 1., .35, 0., True), ('pd_side', 1., .35, 0., True)] results = [] for label, freq, amplitude, force, physics in cases: data = mujoco.MjData(model) data.qpos[:3] = [0, 0, .79] data.qpos[3:7] = [1, 0, 0, 0] data.qpos[qids] = q0 mujoco.mj_forward(model, data) frames, log = [], [] for f in range(240): t = f / 24 ref = q0.copy() for i, j in enumerate(joints): name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_JOINT, j) if 'right_shoulder_pitch' in name: ref[i] -= amplitude * math.sin(2 * math.pi * freq * t / 4) if 'right_elbow' in name: ref[i] += amplitude * .45 * math.sin(2 * math.pi * freq * t / 4) if physics: for step in range(21 if f % 3 != 2 else 20): kp = 35 if label == 'pd_soft' else 80 torque = kp * (ref - data.qpos[qids]) - 3 * data.qvel[vids] data.ctrl[:] = np.clip(torque, model.actuator_ctrlrange[:, 0], model.actuator_ctrlrange[:, 1]) data.xfrc_applied[1, 0] = force if 2 < t < 2.3 else 0 mujoco.mj_step(model, data) else: data.qpos[qids] = ref mujoco.mj_forward(model, data) camera = mujoco.MjvCamera() camera.lookat[:] = [data.qpos[0], data.qpos[1], .65] camera.distance = 2.8 camera.azimuth = 0 if label == 'pd_side' else 140 camera.elevation = -16 renderer.update_scene(data, camera=camera) frames.append(renderer.render().copy()) log.append({'t': round(t, 4), 'joint_rmse_rad': float(np.sqrt(np.mean((ref-data.qpos[qids])**2))), 'base_z_m': float(data.qpos[2]), 'contact_count': int(data.ncon)}) encode(out / f'g1-{label}.mp4', frames) results.append({'id': label, 'physics': physics, 'reference_frequency_multiplier': freq, 'reference_amplitude_rad': amplitude, 'external_push_newton': force, 'samples': log}) print('G1', label, log[-1], flush=True) renderer.close() (out / 'g1-results.json').write_text(json.dumps({'method': 'Synthetic arm reference / plain PD, NOT SuperMimic', 'model': str(model_path), 'model_sha256': hashlib.sha256(model_path.read_bytes()).hexdigest(), 'actuators': names, 'runs': results}, indent=2)) def pendulum(out): from PIL import Image, ImageDraw rng = np.random.default_rng(20261112) xml = '''''' m = mujoco.MjModel.from_xml_string(xml) d = mujoco.MjData(m) rows, targets = [], [] for _ in range(5000): theta, omega, u = rng.uniform(-1.2, 1.2), rng.uniform(-3, 3), rng.uniform(-6, 6) mujoco.mj_resetData(m, d) d.qpos[0], d.qvel[0], d.ctrl[0] = theta, omega, u mujoco.mj_step(m, d) rows.append([theta, omega, math.sin(theta), u, 1]) targets.append([d.qpos[0], d.qvel[0]]) weights = np.linalg.lstsq(rows, targets, rcond=None)[0] def predict(s, u): return np.array([s[0], s[1], np.sin(s[0]), u, 1]) @ weights def rollout(us): states = [np.array([0., 0.])] for u in us: states.append(predict(states[-1], u)) return np.array(states) horizon = 32 us = np.zeros(horizon) losses = [] target = .6 for _ in range(240): states = rollout(us) # Analytic BPTT through a learned one-step dynamics regressor. lam = np.zeros(2) grad = np.zeros(horizon) loss = 0. for k in range(horizon-1, -1, -1): sn = states[k+1] err = sn - np.array([target, 0.]) loss += 8*err[0]**2 + .1*err[1]**2 + .003*us[k]**2 lam += np.array([16*err[0], .2*err[1]]) grad[k] = weights[3] @ lam + .006*us[k] jac = np.stack([weights[0]+np.cos(states[k,0])*weights[2], weights[1]], axis=1) lam = jac.T @ lam losses.append(float(loss)) us = np.clip(us-.015*grad, -6, 6) results = {'method': 'MuJoCo pendulum / fitted differentiable dynamics / analytic BPTT, NOT SuperMimic', 'training_transitions': len(rows), 'horizon': horizon, 'seed': 20261112, 'optimization_losses': losses, 'weights': weights.tolist(), 'runs': []} for label in ['unoptimized', 'optimized', 'model_drift', 'changed_damping']: mujoco.mj_resetData(m, d) if label == 'model_drift': d.qpos[0] = 1.8 if label == 'changed_damping': m.dof_damping[0] = 1.2 else: m.dof_damping[0] = .1 predicted = np.array([float(d.qpos[0]), 0.]) frames, log = [], [] for f in range(288): # Sample-and-hold or receding-horizon replay, clearly an explanation experiment. for substep in range(2): k = min(2*f+substep, horizon-1) u = 0. if label == 'unoptimized' else float(us[k]) if label == 'model_drift': u = 5 * math.sin(f/13) d.ctrl[0] = u mujoco.mj_step(m, d) predicted = predict(predicted, u) im = Image.new('RGB', (960,540), '#0b1726') draw = ImageDraw.Draw(im) for x in range(0,960,80): draw.line((x,0,x,540),fill='#18283b') pivot = np.array([480, 115]) def rod(angle, color, width): tip = pivot+np.array([math.sin(angle), math.cos(angle)])*285 draw.line([tuple(pivot), tuple(tip)], fill=color, width=width) draw.ellipse((tip[0]-14,tip[1]-14,tip[0]+14,tip[1]+14), fill=color) rod(target, '#ffcf65', 3) rod(float(predicted[0]), '#55d8ee', 5) rod(float(d.qpos[0]), '#eef4fa', 13) draw.text((35,35), f'MuJoCo / {label} / t={f/25:.2f}s', fill='white') draw.text((35,480), f'WHITE: physics CYAN: learned prediction GOLD: target | u={u:.2f} Nm', fill='white') frames.append(np.array(im)) log.append({'t':f/25, 'theta_sim':float(d.qpos[0]),'theta_model':float(predicted[0]),'torque':u}) encode(out/f'pendulum-{label}.mp4',frames,25) results['runs'].append({'id':label,'samples':log}) print('PENDULUM',label,flush=True) (out/'pendulum-results.json').write_text(json.dumps(results,indent=2)) if __name__ == '__main__': p = argparse.ArgumentParser() p.add_argument('--out', type=Path, required=True) p.add_argument('--g1-model', type=Path) args = p.parse_args() args.out.mkdir(parents=True, exist_ok=True) (args.out/'execution.json').write_text(json.dumps({'host':platform.node(), 'python':platform.python_version(), 'mujoco':mujoco.__version__, 'physics_device':'CPU', 'renderer':os.getenv('MUJOCO_GL'), 'label':'Independent explanatory experiment; not SuperMimic reproduction'},indent=2)) if not (args.out/'pendulum-results.json').exists(): pendulum(args.out) if args.g1_model: g1(args.out,args.g1_model)