"""Static file server for the bench, plus a small control API.
The bench is a static page, but picking an incentive is only useful if you can
then run it. This adds enough of a backend to launch training and rollout jobs
as subprocesses and report their progress back.
./serve.py [port] default 8901
open http://127.0.0.1:8901/bench/index.html
Binds to loopback only. It runs local subprocesses on request, so do not expose
it beyond this machine.
"""
import collections
import json
import math
import os
import re
import signal
import subprocess
import sys
import threading
import time
import urllib.request
from functools import partial
from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
import mujoco
import numpy as np
ROOT = Path(__file__).parent.resolve()
PY = str(ROOT / ".venv/bin/python")
TRAINER = ROOT / "trainer"
sys.path.insert(0, str(TRAINER))
from tasks import catalogue # noqa: E402
from course import COURSES # noqa: E402
import controller as CTRL # noqa: E402 the ONE implementation of the laws
import firmware as FW # noqa: E402 the laws as COMMANDS - free roam's brain
# A single persistent simulation you can drive by hand. Being able to move one
# joint and watch what the whole robot does is the fastest way to catch a sign
# error, a bad limit, or a leg that pushes the wrong way - the sort of thing
# that silently wrecks a training run.
LIVE = {"env": None, "lock": threading.Lock(), "task": "balance",
"rec": False, "obs": [], "act": []}
# --- Ghrups link: report the sim robot's georeferenced position/heartbeat to Ghrups so it appears
# on the map and shows online, exactly like a person's phone. The device token authenticates as the
# robot; the control-panel URL is where Ghrups opens this bench for that device. Runtime config so
# the same code works on the laptop (dev) and on the always-on server (demo). ---
GHRUPS_API = "https://www.ghrups.com/api/"
GHRUPS_TOKEN_FILE = ROOT / ".ghrups_device_token"
GHRUPS_CONTROL_URL = os.environ.get("GHRUPS_CONTROL_URL", "") # e.g. https://robots.ghrups.com/bench/index.html
# Where the robot idles when it hasn't been driven yet (Portishead), so it always has a position.
GHRUPS_HOME_LAT = float(os.environ.get("GHRUPS_HOME_LAT", "51.4836"))
GHRUPS_HOME_LNG = float(os.environ.get("GHRUPS_HOME_LNG", "-2.7648"))
_GHRUPS = {"last": 0.0, "token": None, "lat": GHRUPS_HOME_LAT, "lng": GHRUPS_HOME_LNG}
def _ghrups_token():
if _GHRUPS["token"] is None:
try:
_GHRUPS["token"] = GHRUPS_TOKEN_FILE.read_text().strip()
except Exception:
_GHRUPS["token"] = ""
return _GHRUPS["token"]
def _ghrups_report(lat, lng):
"""Throttled, non-blocking device_report to Ghrups (position + heartbeat + control feed)."""
tok = _ghrups_token()
if not tok:
return
now = time.time()
if now - _GHRUPS["last"] < 8.0:
return
_GHRUPS["last"] = now
_GHRUPS["lat"], _GHRUPS["lng"] = lat, lng # remember for the idle heartbeat
payload = {"device_token": tok, "lat": lat, "lng": lng, "online_state": "online"}
if GHRUPS_CONTROL_URL:
payload["feeds"] = [{"type": "control", "title": "Control panel", "url": GHRUPS_CONTROL_URL}]
def _post():
try:
req = urllib.request.Request(
GHRUPS_API + "device_report.php",
data=json.dumps(payload).encode(),
headers={"Content-Type": "application/json"})
urllib.request.urlopen(req, timeout=8).read()
except Exception as e:
print("ghrups report failed:", e)
threading.Thread(target=_post, daemon=True).start()
def _ghrups_heartbeat_loop():
"""Keep the robot 'online' in Ghrups even when nobody's driving — a real robot stays connected.
Posts its last known position (or its Portishead home) on a slow cadence."""
while True:
try:
time.sleep(30)
if _ghrups_token():
_ghrups_report(_GHRUPS.get("lat", GHRUPS_HOME_LAT),
_GHRUPS.get("lng", GHRUPS_HOME_LNG))
except Exception as e:
print("ghrups heartbeat error:", e)
def _render_eye(cam="eye_left", w=256, h=192):
"""Render one of the robot's face cameras (mono) from the live sandbox env as JPEG bytes.
Returns None when there's no sandbox robot yet. Rendered under the sim lock so it doesn't race
the stepper. The renderer is cached per model and rebuilt when the env (world) changes."""
sb = LIVE.get("sb") or {}
env = sb.get("env")
if env is None:
return None
with LIVE["lock"]:
try:
rz = LIVE.get("eye_rz")
if rz is None or LIVE.get("eye_model") is not env.model:
if rz is not None:
try:
rz.close()
except Exception:
pass
rz = mujoco.Renderer(env.model, h, w) # needs a GL backend (EGL/OSMesa)
LIVE["eye_rz"] = rz
LIVE["eye_model"] = env.model
rz.update_scene(env.data, camera=cam)
img = rz.render()
except Exception as e:
# No GL backend on this host (headless without EGL/OSMesa) — eyes unavailable, but
# physics/presence/control panel all keep working.
print("eye render unavailable:", e)
return None
import io
from PIL import Image
buf = io.BytesIO()
Image.fromarray(img).save(buf, "JPEG", quality=72)
return buf.getvalue()
# The IK and the trim used to be re-implemented here, and the copies drifted
# from the trainer's (the foot-shift sign was inverted in this file relative
# to the law that was actually scored). One implementation now, in
# trainer/controller.py; these aliases keep the call sites readable.
leg_ik = CTRL.stance_to_leg
lean_trim_deg = CTRL.lean_trim_deg
# The learner runs in a thread so the bench can watch it. It publishes the best
# gains it has found so far, and the live view drives the robot with them, so
# what you see IS the current state of learning rather than a summary of it.
LEARN = {"running": False, "round": 0, "iters": 0, "best": None,
"score": None, "gains": None, "history": [], "stop": False}
try:
_acc = json.load(open(ROOT / "runs" / "balance_gains.json"))
LEARN["accepted"] = _acc
LEARN["gains"] = _acc["gains"] # the bench remembers what you kept
except Exception:
pass
def _polish_thread(stage, run_id, rounds=15, secs=8.0):
sys.path.insert(0, str(TRAINER))
import learn_balance as LB
def on_round(i, n, s, best, g):
if LEARN.get("run_id") != run_id:
return
LEARN.update(round=i, iters=n, score=s, best=best, gains=g,
names=LB.STAGES.get(stage, {}).get("names"),
out_of=100.0)
LEARN["history"].append({"round": i, "score": round(s, 2),
"best": round(best, 2)})
def on_eval(done, total):
if LEARN.get("run_id") == run_id:
LEARN["trial"] = done
LEARN["trials"] = total
def _stop():
return (LEARN.get("stop") or LEARN.get("run_id") != run_id)
try:
res = LB.polish_stage(stage, rounds=rounds, secs=secs,
on_round=on_round, should_stop=_stop,
on_eval=on_eval)
if LEARN.get("run_id") == run_id:
if res and res.get("kept"):
# write the flattened gains into the bake, keeping the
# exam score fields honest (re-measured, not copied)
p = ROOT / "runs" / f"{stage}_gains.json"
rec = json.load(open(p))
rec["gains"] = res["gains"]
rec["score"] = res["raw_after"]
rec["wave_before"] = round(res["wave_before"], 3)
rec["wave_after"] = round(res["wave_after"], 3)
rec["polished"] = True
json.dump(rec, open(p, "w"), indent=1)
LEARN["polish_result"] = {
"kept": True,
"wave_cut_pct": round(100.0 * (
1 - res["wave_after"]
/ max(res["wave_before"], 1e-9)), 1),
"raw_before": round(res["raw_before"], 2),
"raw_after": round(res["raw_after"], 2)}
else:
LEARN["polish_result"] = {"kept": False}
except Exception as e: # noqa: BLE001
LEARN["error"] = f"{type(e).__name__}: {e}"
finally:
if LEARN.get("run_id") == run_id:
LEARN.update(running=False, polish=False)
def _learn_thread(iters, secs, stage="balance", run_id=0, g0=None,
stall=8, dirs=6, nu=0.3):
sys.path.insert(0, str(TRAINER))
import learn_balance as LB
# the operator's stopping rule, automated: repeated rounds with no gain
# means done - stop and wait for the bake click
prog = {"round": 0, "best": None, "best_round": 0}
def on_round(i, n, s, best, g):
prog["round"] = i
if prog["best"] is None or best > prog["best"] + 1e-9:
prog["best"] = best
prog["best_round"] = i
# Only the CURRENT run may write. Stopping a search only takes effect
# between rounds, so an old thread stays alive for a while after the
# next one starts - and both were writing here, which showed one
# stage's parameter names beside another stage's numbers.
if LEARN.get("run_id") != run_id:
return
LEARN.update(round=i, iters=n, score=s, best=best, gains=g,
names=LB.STAGES.get(stage, {}).get("names"),
out_of=LB.stage_max(stage, secs))
LEARN["history"].append({"round": i, "score": round(s, 3),
"best": round(best, 3)})
try:
# Seed from the hand-tuned controller. Starting a search at a known
# good answer is not cheating - it is the only sensible thing to do
# when you have one. It also makes the run a genuine test: anything it
# reports is an improvement ON the thing we already trust, not a
# rediscovery of it.
def _stop():
if LEARN["stop"] or LEARN.get("run_id") != run_id:
return True
if stall and prog["round"] - prog["best_round"] >= stall \
and prog["round"] >= max(6, stall):
LEARN["plateaued"] = True
return True
return False
LEARN.pop("plateaued", None)
def on_eval(k, n):
# trial-by-trial progress WITHIN a round, so a fresh search
# never reads as a hang while it scores its first candidates
if LEARN.get("run_id") == run_id:
LEARN["evals"] = (f"{k}/{n} trials this round" if k
else "scoring the seed")
LB.search_stage(stage, iters=iters, secs=secs, on_round=on_round,
g0=g0, should_stop=_stop, on_eval=on_eval,
dirs=dirs, nu=nu,
# jump explores WIDE: its newest slots seed at
# zero, and at the default 0.05 floor they move
# in 0.015-size steps - the v94 relearn spent
# 300 rounds never touching them (measured).
# Other stages keep the polish floor their
# bakes were learned under.
scale_floor=(0.2 if stage == "jump" else 0.05))
except Exception as e:
# visible in runs/serve.out too - a search that dies mid-run used to
# leave only a silent LEARN["error"] the next start wiped, and the
# partial best got baked as if the search had finished
import traceback
traceback.print_exc()
if LEARN.get("run_id") == run_id:
LEARN["error"] = f"{type(e).__name__}: {e}"
finally:
if LEARN.get("run_id") == run_id:
LEARN["running"] = False
# pushes thrown at the robot DURING this run were pended - bank
# them now so the next run's exam includes them
_bank_pushes(LIVE.pop("push_pending", []))
def _bank_pushes(evs):
"""Append recorded operator pushes to the exam's push bank."""
if not evs:
return
try:
bank = json.load(open(ROOT / "runs" / "push_bank.json")).get("events", [])
except (OSError, ValueError):
bank = []
for ev in evs:
bank.append({"n": round(float(ev["n"]), 1),
"dur": round(float(ev["dur"]), 2)})
json.dump({"events": bank[-10:]},
open(ROOT / "runs" / "push_bank.json", "w"), indent=1)
def mujoco_jid(env, name):
import mujoco
return mujoco.mj_name2id(env.model, mujoco.mjtObj.mjOBJ_JOINT, name)
def mujoco_step(env):
import mujoco
mujoco.mj_step(env.model, env.data)
def _place_on_table(env, stance_mm=420.0):
"""Robot standing on the plate, in pose, plate FLAT - how every round
starts. env.reset() alone leaves the robot at floor height inside the
table geometry, which is where 'leaning whilst flat' came from."""
import mujoco
from assist_lab import stance_to_leg
hp0, kn0 = stance_to_leg(stance_mm / 1000.0)
for n, v in zip(("l_hip", "l_knee", "r_hip", "r_knee"),
(hp0, kn0, hp0, kn0)):
env.data.qpos[env.model.jnt_qposadr[mujoco_jid(env, n)]] = v
env.data.qpos[0:3] = [0.0, 0.0, 0.73 + stance_mm / 1000.0 + 0.003]
env.data.qpos[3:7] = [1.0, 0.0, 0.0, 0.0]
env.data.qpos[-1] = 0.0
env.data.qvel[:] = 0.0
env.data.ctrl[-1] = 0.0
mujoco.mj_forward(env.model, env.data)
def _live_env(reset=False, task=None, tether=None, bank_deg=None, world=None):
"""The hand-driven sim.
Tethering is a real joint set, not a per-step overwrite of the body pose.
Forcing qpos every step injects momentum the physics never agreed to, and
the robot creeps backwards whichever way you drive it. A rig built from
slide joints simply cannot topple, and everything else stays honest.
"""
sys.path.insert(0, str(TRAINER))
from env import RobotEnv
task = task or LIVE["task"]
tether = LIVE.get("tether", "free") if tether is None else tether
bank = LIVE.get("bank", 0.0) if bank_deg is None else float(bank_deg)
world_now = LIVE.get("world", "floor") if world is None else world
# A tilted floor cannot be set at runtime - MuJoCo precomputes the frame of
# static world geoms - so changing the gradient rebuilds the model.
if (LIVE["env"] is None or reset or task != LIVE["task"]
or tether != LIVE.get("tether") or bank != LIVE.get("bank")
or world_now != LIVE.get("world")):
LIVE["bank"] = bank
LIVE["world"] = world_now
LIVE["task"], LIVE["tether"] = task, tether
spec_path = str(ROOT / "robots/wheeled_biped.json")
if bank:
import json as _json, copy as _copy, mjcf as _mjcf
_s = _mjcf.load_spec(spec_path); _s.pop("_path", None)
_s.setdefault("world", {}).setdefault("floor", {})["bank_deg"] = bank
spec_path = str(ROOT / "runs" / "_live_banked.json")
(ROOT / "runs").mkdir(exist_ok=True)
_json.dump(_s, open(spec_path, "w"))
env = RobotEnv(spec_path, seed=0, task=task, randomise=False)
if tether != "free":
import re as _re, mujoco, mjcf
joints = {
# Hung from a hook at the centre of the chassis: a ball joint
# pins the position but leaves all three rotations free, so the
# robot can lean, topple and turn exactly as it would on a
# string. Constraining translation instead of rotation was the
# mistake - it stopped the reaction that turns the body.
# A slack line: free to slide down its own vertical, so the
# robot settles onto its wheels and they carry the full weight.
# The line only takes over if it drops past standing height -
# it can lean and spin but never hit the floor or run away.
# limited slide = the string has a length: the robot stands on
# its wheels with full weight (the line slack), can lean and
# spin, and if it topples the line catches it after ~60mm of
# sag instead of letting it dangle.
"hook": ''
'',
"upright": ''
''
''
'',
"pinned": '',
}[tether]
xml = mjcf.spec_to_mjcf(env.spec, props=env.props)
xml = xml.replace('', joints)
env.model = mujoco.MjModel.from_xml_string(xml)
env.data = mujoco.MjData(env.model)
env._rebind()
if world_now == "table":
# the tilting table the level stage is scored on - the demo must
# show the same world the search lives in, or what you watch is
# a different robot (the stale slew demo ran the new gains with
# the OLD sign convention: wrong leg up, leaning on flat)
import mujoco, mjcf as _mjcf2
import learn_balance as _LB2
xml = _mjcf2.spec_to_mjcf(env.spec, props=env.props)
xml = xml.replace("", _LB2.PLATFORM_XML + "\n ")
xml = xml.replace("", _LB2.PLATFORM_ACT + "\n ")
env.model = mujoco.MjModel.from_xml_string(xml)
env.data = mujoco.MjData(env.model)
env._rebind()
n = len(env.spec["actuators"])
env.ctrl_lo = env.ctrl_lo[:n]
env.ctrl_hi = env.ctrl_hi[:n]
env.tau_limit = env.tau_limit[:n]
if world_now == "bumps":
# the threshold ridges the bump stage is scored on - props, so
# the physics AND the bench renderer both carry them
import mujoco, mjcf as _mjcf3
import learn_balance as _LB3
env.props = list(env.props or []) + _LB3.bump_props(0.010)
xml = _mjcf3.spec_to_mjcf(env.spec, props=env.props)
env.model = mujoco.MjModel.from_xml_string(xml)
env.data = mujoco.MjData(env.model)
env._rebind()
LIVE["env"] = env
LIVE["obs_now"] = env.reset()
if world_now == "table":
# reset spawns at floor height - UNDER the slab. The on-fall
# guard used to catch this by accident (the old zeroed-qpos
# reset exploded and triggered a reposition); a clean reset just
# leaves the robot standing politely beneath the table.
_place_on_table(env)
env._cache_truth()
# The hand-driving controller keeps integrator state - odometry, the
# ramped speed command, the leg targets, which knee branch it is in.
# None of that means anything for a fresh robot, and carrying it over
# left odometry sitting at -0.8m straight after a reset, biasing the
# lean target before the robot had moved at all.
LIVE.pop("assist", None)
LIVE["odo"] = 0.0
return LIVE["env"]
JOBS = {}
JOB_LOCK = threading.Lock()
_next_id = [0]
SAFE_NAME = re.compile(r"^[A-Za-z0-9._-]+$")
# Rolling black box of live steps. 60k entries is about 30 minutes at 33Hz,
# so anything you notice and report is still in the buffer when we look.
BLACKBOX = collections.deque(maxlen=60000)
STATE = ROOT / "runs" / "jobs.json"
CORES = os.cpu_count() or 4
def _alive(pid):
try:
os.kill(pid, 0)
return True
except OSError:
return False
def _save_jobs():
"""Jobs are detached subprocesses, so a server restart should not lose
sight of them. Persist the bookkeeping and re-adopt by pid on startup."""
try:
STATE.parent.mkdir(parents=True, exist_ok=True)
with open(STATE, "w") as f:
json.dump([{k: v for k, v in j.items() if k != "proc"}
for j in JOBS.values()], f)
except OSError:
pass
def _load_jobs():
if not STATE.exists():
return
try:
rows = json.load(open(STATE))
except (OSError, ValueError):
return
# Keep only jobs still running: finished ones from previous sessions are
# noise, and their pids get recycled, which makes "is it alive" unreliable.
rows = [r for r in rows if _alive(r.get("pid", -1))]
for r in rows:
r["proc"] = None # adopted: no handle, poll by pid
r["adopted"] = True
JOBS[r["id"]] = r
n = int(re.sub(r"\D", "", r["id"]) or 0)
_next_id[0] = max(_next_id[0], n)
def _running_workers():
total = 0
for j in JOBS.values():
if j["kind"] != "train" or not _job_alive(j):
continue
argv = j["argv"]
if "--workers" in argv:
total += int(argv[argv.index("--workers") + 1])
return total
def _job_alive(j):
return j["proc"].poll() is None if j.get("proc") else _alive(j["pid"])
# Training saturates the box. Two runs on four cores do not go twice as fast,
# they each go a third the speed and starve the other, so only one heavy job
# is allowed to exist at a time. Rollouts are a second of CPU and are merely
# de-prioritised rather than blocked.
HEAVY = {"train", "course"}
def _running(kinds):
return [j for j in JOBS.values() if j["kind"] in kinds and _job_alive(j)]
def _busy():
r = _running(HEAVY)
return r[0] if r else None
def _launch(kind, argv, run_dir=None, label=""):
with JOB_LOCK:
_next_id[0] += 1
jid = f"job{_next_id[0]}"
log_path = (run_dir or ROOT / "runs") / f"{jid}.out"
log_path.parent.mkdir(parents=True, exist_ok=True)
fh = open(log_path, "w")
# Rollouts run nice so that watching a replay never slows a training run.
pre = (lambda: os.nice(10)) if kind == "rollout" else None
proc = subprocess.Popen([PY] + argv, cwd=str(TRAINER), stdout=fh,
stderr=subprocess.STDOUT, start_new_session=True,
preexec_fn=pre)
JOBS[jid] = {
"id": jid, "kind": kind, "label": label, "argv": argv,
"pid": proc.pid, "proc": proc, "log": str(log_path),
"run_dir": str(run_dir) if run_dir else None,
"started": time.time(),
}
_save_jobs()
return JOBS[jid]
def _job_view(j):
proc = j.get("proc")
if proc is not None:
code = proc.poll()
status = "running" if code is None else ("done" if code == 0 else f"failed ({code})")
else:
# Adopted across a server restart: all we can check is whether the pid
# is still there, so a finished job reports "ended" not "done".
status = "running" if _alive(j["pid"]) else "ended"
out = {k: j.get(k) for k in ("id", "kind", "label", "pid", "started", "run_dir")}
out["elapsed_s"] = int(time.time() - j["started"])
out["status"] = status
tail = ""
try:
with open(j["log"]) as f:
lines = [l.rstrip() for l in f.readlines() if l.strip()]
tail = lines[-1] if lines else ""
except OSError:
pass
out["tail"] = tail[-300:]
# Training progress comes from the structured log, not the console text.
if j["run_dir"]:
lp = Path(j["run_dir"]) / "log.jsonl"
if lp.exists():
try:
with open(lp) as f:
last = f.readlines()[-1]
out["progress"] = json.loads(last)
except (OSError, ValueError, IndexError):
pass
return out
class Handler(SimpleHTTPRequestHandler):
def end_headers(self):
# Never let the browser cache the bench. Every edit to a js file should
# be live on reload; a stale cached module looks exactly like "the
# change did not work".
self.send_header("Cache-Control", "no-store, no-cache, must-revalidate")
self.send_header("Pragma", "no-cache")
super().end_headers()
def _send(self, obj, code=200):
body = json.dumps(obj).encode()
self.send_response(code)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(body)))
self.send_header("Cache-Control", "no-store")
self.end_headers()
self.wfile.write(body)
def _query(self, key):
from urllib.parse import urlparse, parse_qs
return (parse_qs(urlparse(self.path).query).get(key) or [None])[0]
def _body(self):
n = int(self.headers.get("Content-Length") or 0)
return json.loads(self.rfile.read(n) or b"{}")
# -- GET ---------------------------------------------------------------
def do_GET(self):
if self.path.startswith("/api/eye"):
# Mono camera feed: render one of the robot's face cameras (what the robot SEES) as a
# JPEG. The control panel swaps its 3D sensed-data scene for this.
from urllib.parse import urlparse, parse_qs
q = parse_qs(urlparse(self.path).query)
cam = (q.get("cam") or ["eye_left"])[0]
jpg = _render_eye(cam)
if jpg is None:
self.send_response(503); self.end_headers(); return
self.send_response(200)
self.send_header("Content-Type", "image/jpeg")
self.send_header("Cache-Control", "no-store")
self.send_header("Content-Length", str(len(jpg)))
self.end_headers()
self.wfile.write(jpg)
return
if self.path.startswith("/api/incentives"):
specs = sorted(p.name for p in (ROOT / "robots").glob("*.json"))
return self._send({"incentives": catalogue(), "specs": specs})
if self.path.startswith("/api/jobs"):
b = _busy()
return self._send({
"jobs": [_job_view(j) for j in JOBS.values()],
"busy": bool(b),
"busy_label": b["label"] if b else None,
"cores": CORES,
})
if self.path.startswith("/api/runs"):
runs = []
for d in sorted((ROOT / "runs").glob("*/")):
best = d / "policy_best.json"
if not best.exists():
continue
try:
meta = json.load(open(best)).get("meta", {})
except (OSError, ValueError):
meta = {}
runs.append({
"name": d.name,
"task": meta.get("task", "?"),
"steps": meta.get("steps", 0),
"ep_len": meta.get("ep_len", 0),
"policy": f"../runs/{d.name}/policy_best.json",
"mtime": best.stat().st_mtime,
})
runs.sort(key=lambda r: -r["mtime"])
return self._send({"runs": runs})
if self.path.startswith("/api/courses"):
return self._send({"courses": [
{"name": n, "label": c["label"], "description": c["description"],
"stages": [{"task": s["task"], "steps": s["steps"], "why": s["why"]}
for s in c["stages"]]}
for n, c in COURSES.items()]})
if self.path.startswith("/api/course"):
run = self._query("run")
if not run:
# No run named: report the most recently started course.
cands = sorted((ROOT / "runs").glob("*/course.json"),
key=lambda p: -p.stat().st_mtime)
if not cands:
return self._send({"course": None})
return self._send({"course": json.load(open(cands[0]))})
if not SAFE_NAME.match(run):
return self._send({"error": "bad run name"}, 400)
p = ROOT / "runs" / run / "course.json"
if not p.exists():
return self._send({"course": None})
return self._send({"course": json.load(open(p))})
if self.path.startswith("/api/progress"):
run = self._query("run") or "johnny6"
course = self._query("course") or "standard"
if not SAFE_NAME.match(run):
return self._send({"error": "bad run name"}, 400)
cp = ROOT / "runs" / run / "course.json"
state = json.load(open(cp)) if cp.exists() else None
# A killed stage never gets to write its own ending, so the file
# still claims it is training. Trusting it leaves the UI stuck on
# "learning" forever. Reconcile against what is actually alive.
alive = bool(_busy())
if state and not alive:
for st in state.get("stages", []):
if st.get("status") in ("training", "examining"):
st["status"] = "stopped"
if state.get("status") == "running":
state["status"] = "idle"
stages = []
plan = COURSES[course]["stages"]
for i, st in enumerate(plan):
cur = (state or {}).get("stages", [{}]*len(plan))[i] if state else {}
sr = cur.get("run") or f"{run}__{st['task']}"
row = {
"index": i + 1, "task": st["task"], "why": st["why"],
"resumable": (ROOT / "runs" / sr / "policy_latest.json").exists(),
"steps": cur.get("steps", st["steps"]),
"status": cur.get("status", "pending"),
"exam": cur.get("exam"), "run": sr,
"gate_pass": st["gate"].get("pass_rate"),
"gate_frac": st["gate"].get("survive_frac"),
}
# live progress for whichever stage is training right now
lp = ROOT / "runs" / sr / "log.jsonl"
if lp.exists():
try:
with open(lp) as f:
lines = f.readlines()
row["live"] = json.loads(lines[-1])
except (OSError, ValueError, IndexError):
pass
snaps = sorted((ROOT / "runs" / sr / "snapshots").glob("*.json")) \
if (ROOT / "runs" / sr / "snapshots").exists() else []
rows_s = []
for x in snaps:
try:
sm = json.load(open(x))["summary"]
except (OSError, ValueError, KeyError):
sm = {}
rows_s.append({"steps": int(x.stem),
"url": f"../runs/{sr}/snapshots/{x.name}",
"survived_s": sm.get("survived_s", 0),
"fell": sm.get("fell", True)})
row["snapshots"] = rows_s
row["policy"] = (f"../runs/{sr}/policy_best.json"
if (ROOT / "runs" / sr / "policy_best.json").exists()
else None)
stages.append(row)
b = _busy()
return self._send({"run": run, "course": course, "stages": stages,
"busy": bool(b),
"busy_label": b["label"] if b else None,
"cores": CORES})
if self.path.startswith("/api/snapshots"):
run = self._query("run") or ""
if not SAFE_NAME.match(run):
return self._send({"error": "bad run name"}, 400)
d = ROOT / "runs" / run / "snapshots"
snaps = []
for p in sorted(d.glob("*.json")):
try:
s = json.load(open(p))["summary"]
except (OSError, ValueError, KeyError):
continue
snaps.append({
"steps": int(p.stem),
"url": f"../runs/{run}/snapshots/{p.name}",
"survived_s": s["survived_s"],
"fell": s["fell"],
"success": s.get("success", False),
})
return self._send({"snapshots": snaps})
if self.path.startswith("/api/learn"):
# Which stages have been signed off by the search. The course tracks
# its own PPO status, so without this it still thought balance was
# unlearned and refused to let you start the next step.
done = {}
for f in (ROOT / "runs").glob("*_gains.json"):
try:
r = json.load(open(f))
done[r.get("stage", f.stem.replace("_gains", ""))] = r
except (OSError, ValueError):
pass
import learn_balance as _LBl
return self._send({k: v for k, v in LEARN.items() if k != "history"}
| {"history": LEARN["history"][-120:],
"accepted_stages": done,
# the ladder lives in ONE place (trainer);
# the UI renders whatever this says, so
# adding a stage never needs a JS edit
"ladder": _LBl.STAGE_ORDER})
if self.path.startswith("/api/blackbox"):
# Dump the rolling buffer so a one-off glitch can be examined after
# the fact. ?mark=... just drops a labelled bookmark in the stream.
q = {}
if "?" in self.path:
for kv in self.path.split("?", 1)[1].split("&"):
k, _, v = kv.partition("=")
q[k] = v
if "mark" in q:
BLACKBOX.append({"mark": q["mark"], "wall": round(time.time(), 3)})
return self._send({"marked": q["mark"], "steps": len(BLACKBOX)})
rows = list(BLACKBOX)[-int(q.get("n", 4000)):]
out = ROOT / "runs" / "blackbox.json"
out.parent.mkdir(parents=True, exist_ok=True)
with open(out, "w") as f:
json.dump(rows, f)
return self._send({"steps": len(rows), "total": len(BLACKBOX),
"file": str(out)})
if self.path.startswith("/api/replays"):
d = ROOT / "bench/replays"
files = sorted(d.glob("*.json"), key=lambda p: -p.stat().st_mtime)
return self._send({"replays": [f"replays/{p.name}" for p in files]})
return super().do_GET()
# -- POST --------------------------------------------------------------
def do_POST(self):
try:
if self.path.startswith("/api/train"):
return self._train(self._body())
if self.path.startswith("/api/rollout"):
return self._rollout(self._body())
if self.path.startswith("/api/step"):
return self._step(self._body())
if self.path.startswith("/api/course"):
return self._course(self._body())
if self.path.startswith("/api/clone"):
return self._clone(self._body())
if self.path.startswith("/api/learn"):
b = self._body()
if b.get("stop"):
LEARN["stop"] = True
return self._send({"stopping": True})
if b.get("load"):
# Put a saved stage back in charge, without training
# anything. The live view runs whatever gains are current,
# so this is how you go back to a result you trust.
# NEVER silently over a live search: a load retires the
# running thread (run_id bump below), and one stray click
# on a baked row killed the crouch bench's first run
# while the operator watched round 0 - stop first, or
# pass force.
if LEARN.get("running") and not b.get("force"):
return self._send({"error":
f"a search is running ({LEARN.get('stage')}) - "
f"stop it before loading a saved stage"}, 409)
try:
rec = json.load(open(ROOT / "runs" / f"{b['load']}_gains.json"))
except (OSError, ValueError):
return self._send({"error": f"nothing saved for {b['load']}"}, 404)
# The gate check_stages runs by hand, enforced here: a
# bake whose parameter names no longer match the stage's
# layout scatters numbers into the wrong slots (stop's
# brake gain once landed in recover's brace slot, a 25x
# overdose). Refuse it rather than run it.
import learn_balance as _LBg
_stage_l = rec.get("stage", b["load"])
_want = _LBg.STAGES.get(_stage_l, {}).get("names")
if _want and rec.get("names") != _want:
return self._send({"error":
f"{_stage_l} bake is STALE (layout changed) - "
f"relearn it rather than load it"}, 409)
if rec.get("score_version") not in (None, _LBg.SCORE_VERSION):
# gains are still usable, but the recorded score is
# from a different objective - do not display it as
# comparable
rec = dict(rec)
rec["score"] = None
rec["stale_score"] = True
LEARN["stop"] = True
# Retire any run outright: bumping run_id makes the orphan
# thread's writes dead on arrival. Without this, loading a
# stage while a search was mid-round left the search
# overwriting the loaded gains every round - the demo ran
# one stage's candidates under another stage's name, and
# its controls (the bob gate slider) appeared dead.
LEARN["run_id"] = LEARN.get("run_id", 0) + 1
LEARN["running"] = False
LEARN["gains"] = rec["gains"]
LEARN["names"] = rec.get("names")
LEARN["stage"] = rec.get("stage", b["load"])
LEARN["accepted"] = rec
LEARN["best"] = rec.get("score")
return self._send({"loaded": rec})
if b.get("accept"):
# Keep what it found. Saved as readable gains with the
# score and the date, because these are numbers you would
# type into the robot, not weights to be reloaded blindly.
g = b.get("gains") or LEARN.get("gains")
if not g:
return self._send({"error": "nothing to accept"}, 409)
LEARN["stop"] = True
stage = LEARN.get("stage", "balance")
out = ROOT / "runs" / f"{stage}_gains.json"
out.parent.mkdir(parents=True, exist_ok=True)
import learn_balance as _LBa
rec = {"stage": stage,
# layout and scorer version are recorded ALWAYS so
# check_stages and the load gate can tell a usable
# bake from a stale one (a bake was once written
# with score null and no layout, which defeated
# every check downstream)
"names": (LEARN.get("names")
or _LBa.STAGES.get(stage, {}).get("names")
or ["kp", "kd", "station", "clamp"]),
"gains": [float(x) for x in g],
"score": LEARN.get("best"),
"unscored": LEARN.get("best") is None,
"score_version": _LBa.SCORE_VERSION,
"rounds": LEARN.get("round"),
"out_of": LEARN.get("out_of"),
"hand_written_scores": 4.80,
"saved": time.strftime("%Y-%m-%d %H:%M:%S")}
# Refuse to silently replace a better result baked under
# the SAME scorer - overwriting a good bake with a worse
# run is how progress quietly walks backwards. force=true
# (the UI asks first) overrides.
try:
old = json.load(open(out))
except (OSError, ValueError):
old = None
if (old and not b.get("force")
and old.get("score") is not None
and rec["score"] is not None
and old.get("score_version") == rec["score_version"]
and old["score"] > rec["score"]):
return self._send({"error":
f"saved {stage} scores {old['score']:.2f}, this "
f"run's best is {rec['score']:.2f}"}, 409)
json.dump(rec, open(out, "w"), indent=2)
LEARN["accepted"] = rec
LEARN["gains"] = rec["gains"] # keep driving with them
# a fresh bench bake changes bob's interpolation - drop
# the demo's cached points so it picks the new law up
LIVE.get("gainstate", {}).pop("bpts", None)
order = _LBa.STAGE_ORDER
nxt = (order[order.index(stage) + 1]
if stage in order and order.index(stage) + 1 < len(order)
else None)
return self._send({"accepted": rec, "file": str(out),
"next": nxt})
if (LEARN["running"] and not b.get("g0")
and LEARN.get("stage") == b.get("stage", "balance")):
# same stage already going: keep it - UNLESS an explicit
# seed came with the request, which always means "restart
# from here" (a reseed was silently swallowed by this
# guard while reporting the old run's numbers).
# Teaching the SAME stage again also revives it: a stop
# request followed by a re-teach used to leave the stop
# flag armed, and the run quietly died at its next round
# boundary while this reply said "already going".
LEARN["stop"] = False
return self._send({"already": True, "stage": LEARN["stage"]})
if LEARN["running"]:
# Preempt. A run only checks its stop flag between rounds,
# and a round is up to a minute - so "stop, wait 250ms,
# start" always hit "already learning", the UI swallowed
# the 409, and clicking Teach on one stage silently left
# the previous stage running. The run_id guard means the
# old thread's writes are ignored from this moment, so
# starting the new stage NOW is safe; the orphan retires
# at its next round boundary.
LEARN["stop"] = True
# Continuation. Same stage under the same scoring resumes from
# the previous best; an explicit g0 carries a working robot
# across a deliberate score change. Without this, five
# relaunches in an hour each silently discarded a learned best
# and made progress look like regression. (An earlier attempt
# at this feature failed its text match and never landed -
# while its commit message said it had. Verified by behaviour
# this time: the new run's first gains must equal the seed.)
import learn_balance as _LB
_g0 = None
if b.get("g0"):
_g0 = [float(x) for x in b["g0"]]
elif (LEARN.get("stage") == b.get("stage", "balance")
and LEARN.get("score_v") == _LB.SCORE_VERSION
and LEARN.get("gains")):
_g0 = list(LEARN["gains"])
if b.get("polish"):
# THE FLATTENER (operator's design): hold the exam
# score, minimize the wave height, from the bake
LEARN.update(running=True, stop=False, round=0,
best=None, score=None, gains=None,
history=[], polish=True)
LEARN.pop("error", None)
LEARN["run_id"] = LEARN.get("run_id", 0) + 1
LEARN["stage"] = b.get("stage", "balance")
LEARN["score_v"] = _LB.SCORE_VERSION
threading.Thread(
target=_polish_thread, daemon=True,
args=(b.get("stage", "balance"),
LEARN["run_id"],
int(b.get("rounds", 15)),
float(b.get("secs", 8.0)))).start()
return self._send({"polishing": True,
"stage": LEARN["stage"]})
LEARN.update(running=True, stop=False, round=0, best=None,
score=None, gains=None, history=[])
LEARN.pop("error", None)
LEARN["run_id"] = LEARN.get("run_id", 0) + 1
LEARN["stage"] = b.get("stage", "balance")
LEARN["score_v"] = _LB.SCORE_VERSION
LEARN["gains"] = None # a new stage starts with no result
LEARN["names"] = None
LEARN.pop("accepted", None) # a new stage is not yet signed off
# jump-sized defaults: 27 slots need hundreds of wide
# rounds, and a strong seed guarantees quiet early
# rounds - the stall=8 rule (tuned for the 4-gain
# benches) killed a jump run at round 9 with every new
# slot still untouched (measured, 16 Aug). The UI can
# still override any of these explicitly.
_big = LEARN["stage"] == "jump"
# jump: episodes must outlive the 10s landing hold
# (v100), and the messy-entry exam is cliff-heavy -
# one failed entry in eight is -1 - so steps are
# GENTLER (nu 0.2), not just wider: 2000 nu-0.3
# candidates found nothing (measured, 16 Aug)
_secs = float(b.get("secs", 8.0))
if _big:
_secs = max(_secs, 12.0)
threading.Thread(target=_learn_thread, daemon=True,
args=(int(b.get("iters",
300 if _big else 40)),
_secs,
LEARN["stage"], LEARN["run_id"],
_g0, int(b.get("stall",
60 if _big else 8)),
int(b.get("dirs",
16 if _big else 6)),
float(b.get(
"nu",
0.2 if _big else 0.3)))).start()
return self._send({"started": True, "stage": LEARN["stage"]})
if self.path.startswith("/api/sandbox"):
return self._sandbox(self._body())
if self.path.startswith("/api/terrain"):
sb = LIVE.get("sb") or {}
envt = sb.get("env")
if envt is not None and envt.model.nhfield:
mt = envt.model
return self._send({
"nrow": int(mt.hfield_nrow[0]),
"ncol": int(mt.hfield_ncol[0]),
"size": [float(v) for v in mt.hfield_size[0]],
"data": [round(float(v), 4)
for v in mt.hfield_data]})
return self._send({"nrow": 0})
if self.path.startswith("/api/live"):
return self._live(self._body())
if self.path.startswith("/api/reset_step"):
return self._reset_step(self._body())
if self.path.startswith("/api/stop"):
return self._stop(self._body())
except Exception as e: # noqa: BLE001
import traceback
traceback.print_exc() # the message alone is not enough to fix it
return self._send({"error": f"{type(e).__name__}: {e}"}, 400)
self._send({"error": "unknown endpoint"}, 404)
def _spec_arg(self, b):
spec = b.get("spec") or "wheeled_biped.json"
if not SAFE_NAME.match(spec):
raise ValueError("bad spec name")
return str(ROOT / "robots" / spec)
def _train(self, b):
task = b.get("task", "drive")
if not SAFE_NAME.match(task):
raise ValueError("bad task name")
steps = max(10_000, min(50_000_000, int(b.get("steps", 3_000_000))))
run = b.get("run") or f"{task}_{time.strftime('%H%M%S')}"
if not SAFE_NAME.match(run):
raise ValueError("bad run name")
workers = int(b.get("workers", 3))
argv = ["train.py", "--task", task, "--steps", str(steps),
"--run", run, "--spec", self._spec_arg(b),
"--envs", str(int(b.get("envs", 24))),
"--workers", str(workers)]
if b.get("no_dr"):
argv.append("--no-dr")
b = _busy()
if b and not b.get("_stale"):
return self._send({
"error": f"{b['label']} is already running. Only one training job "
f"at a time on {CORES} cores - stop it first.",
"busy": _job_view(b)}, 409)
j = _launch("train", argv, run_dir=ROOT / "runs" / run,
label=f"train {task} ({steps/1e6:.1f}M)")
self._send(_job_view(j))
def _rollout(self, b):
out_name = b.get("out") or f"{b.get('task', 'drive')}_latest.json"
if not SAFE_NAME.match(out_name):
raise ValueError("bad output name")
argv = ["rollout.py", "--out", str(ROOT / "bench/replays" / out_name),
"--seconds", str(float(b.get("seconds", 12))),
"--seed", str(int(b.get("seed", 123)))]
if b.get("task"):
argv += ["--task", b["task"]]
if b.get("policy"):
# The bench sends paths as the page sees them ("../runs/x/y.json").
# Resolve against the project root and refuse anything that escapes.
rel = b["policy"].lstrip("/")
while rel.startswith("../"):
rel = rel[3:]
pol = (ROOT / rel).resolve()
if not str(pol).startswith(str(ROOT) + os.sep):
raise ValueError("policy path outside project")
if not pol.exists():
raise ValueError(f"no such policy: {rel}")
argv += ["--policy", str(pol)]
else:
argv += ["--spec", self._spec_arg(b)]
if b.get("no_dr"):
argv.append("--no-dr")
if _running({"rollout"}):
return self._send({"error": "a rollout is already running"}, 409)
j = _launch("rollout", argv,
label=f"rollout {b.get('task', '')} -> replays/{out_name}")
j["replay"] = f"replays/{out_name}"
v = _job_view(j)
v["replay"] = j["replay"]
self._send(v)
def _course(self, b):
name = b.get("course", "standard")
if name not in COURSES:
raise ValueError(f"unknown course {name!r}")
run = b.get("run") or f"course_{name}_{time.strftime('%H%M%S')}"
if not SAFE_NAME.match(run):
raise ValueError("bad run name")
scale = float(b.get("steps_scale", 1.0))
workers = int(b.get("workers", 3))
argv = ["course.py", "--course", name, "--run", run,
"--spec", self._spec_arg(b), "--envs", str(int(b.get("envs", 24))),
"--workers", str(workers), "--steps-scale", str(scale),
"--snapshot-every", str(int(b.get("snapshot_every", 500_000)))]
b = _busy()
if b:
return self._send({
"error": f"{b['label']} is already running. Only one training job "
f"at a time on {CORES} cores - stop it first.",
"busy": _job_view(b)}, 409)
j = _launch("course", argv, run_dir=ROOT / "runs" / run,
label=f"course {name}")
v = _job_view(j)
v["run"] = run
self._send(v)
def _step(self, b):
"""Train exactly one stage of the course, seeded from the one before."""
name = b.get("course", "standard")
if name not in COURSES:
raise ValueError(f"unknown course {name!r}")
run = b.get("run") or "johnny6"
if not SAFE_NAME.match(run):
raise ValueError("bad run name")
idx = int(b.get("index", 1))
if not 1 <= idx <= len(COURSES[name]["stages"]):
raise ValueError("step out of range")
busy = _busy()
if busy:
return self._send({"error": f"{busy['label']} is already running."}, 409)
task = COURSES[name]["stages"][idx-1]["task"]
argv = ["course.py", "--course", name, "--run", run, "--only", str(idx),
"--spec", self._spec_arg(b), "--envs", str(int(b.get("envs", 24))),
"--workers", str(int(b.get("workers", 3))),
"--steps-scale", str(float(b.get("steps_scale", 1.0))),
"--snapshot-every", str(int(b.get("snapshot_every", 250_000)))]
j = _launch("course", argv, run_dir=ROOT / "runs" / run,
label=f"step {idx}: {task}")
v = _job_view(j); v["run"] = run
self._send(v)
def _sandbox(self, b):
# concurrent requests were stepping the same mjData from two
# threads - MuJoCo segfaulted in its collision tree and took the
# process down (the crash report shows mj_step live on two
# threads at once). One body, one stepper, ever.
with LIVE["lock"]:
return self._sandbox_locked(b)
def _sandbox_locked(self, b):
# -------------------------------------------------- THE SANDBOX
# Completely OUTSIDE the learning machinery (operator's spec):
# its own environment, its own state, only the baked firmware -
# bench gains interpolated at the MEASURED hip height, the trim
# curve, the speed loop. No stage logic, no LEARN gains, and NO
# AUTO-RESETS EVER: falls lie where they fell, power-on wakes
# where it lies, and standing back up is the robot's problem (or
# the operator's, by the manual sliders). The only reset is the
# explicit button.
if str(TRAINER) not in sys.path:
sys.path.insert(0, str(TRAINER))
import learn_balance as LB
sb = LIVE.setdefault("sb", {})
sb["reqn"] = sb.get("reqn", 0) + 1
# ONE PILOT AT A TIME (measured: a bench tab idling in free
# roam streamed its zero stick 5x/s underneath an automated
# test's full-left - every "release" the firmware honoured was
# real, just not from the pilot who thought they were flying;
# the same interleave also stepped the physics twice as fast).
# A client that sends a pilot id and claims (or acts) becomes
# THE pilot; everyone else is an observer - read-only, no
# stepping. A pilot silent for 3s lapses so old tabs recover.
pid = b.get("pilot")
now_t = time.time()
if sb.get("pilot") is not None and (
now_t - sb.get("pilot_t", 0.0) > 3.0):
sb["pilot"] = None
if pid is not None:
if (b.get("claim") or b.get("hard_reset") or b.get("wake")
or sb.get("pilot") is None):
sb["pilot"] = pid
if sb.get("pilot") == pid:
sb["pilot_t"] = now_t
else:
b = {"substeps": 0}
elif sb.get("pilot") is not None:
b = {"substeps": 0}
env = sb.get("env")
# THE PLAYGROUND TERRAINS: switching worlds rebuilds the sandbox
# into the chosen landscape, starting collapsed as always
want_world = b.get("world") or sb.get("world") or "flat"
if sb.get("world") and want_world != sb.get("world"):
b = dict(b)
b["hard_reset"] = True
if env is None or b.get("hard_reset"):
if sb.get("off") and sb.get("gain0") is not None and env:
env.model.actuator_gainprm[:] = sb["gain0"]
env.model.actuator_biasprm[:] = sb["bias0"]
from env import RobotEnv
if want_world == "hills":
import copy as _copy
_hspec = json.load(open(
ROOT / "robots" / "wheeled_biped.json"))
_hspec.setdefault("world", {}).setdefault(
"floor", {})["hills"] = True
_hpath = ROOT / "runs" / "_hills_world.json"
json.dump(_hspec, open(_hpath, "w"))
env = RobotEnv(str(_hpath), seed=0, task="bob",
randomise=False)
else:
env = RobotEnv(str(ROOT / "robots" / "wheeled_biped.json"),
seed=0, task="bob", randomise=False)
env.reset()
env.min_h = 0.02
env.min_up = -2.0
# reset delivers the OFF state (operator's spec): the true
# rest - chassis down, knees on their rollers in front,
# wheels behind - and UNPOWERED; waking is a deliberate act
LB.place_sitting(env)
LB.limp_settle(env, 1.2)
sb.clear()
sb.update(env=env, world=want_world,
fw=FW.Firmware(env.control_dt, ROOT / "runs"))
sb["off"] = True
sb["gain0"] = env.model.actuator_gainprm.copy()
sb["bias0"] = env.model.actuator_biasprm.copy()
env.model.actuator_gainprm[:, 0] = 0.0
env.model.actuator_biasprm[:, 1:3] = 0.0
m, d = env.model, env.data
if "push_bid" not in sb:
sb["push_bid"] = mujoco.mj_name2id(
m, mujoco.mjtObj.mjOBJ_BODY, "chassis")
sb["hipj"] = [mujoco_jid(env, n) for n in ("l_hip", "r_hip")]
sb["floor_g"] = mujoco.mj_name2id(
m, mujoco.mjtObj.mjOBJ_GEOM, "floor")
_wb = [mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_BODY, n)
for n in ("l_wheel", "r_wheel")]
sb["tyre_g"] = {gi for gi in range(m.ngeom)
if int(m.geom_bodyid[gi]) in _wb}
sb["wd"] = [m.jnt_dofadr[mujoco_jid(env, n)]
for n in ("l_wheel_j", "r_wheel_j")]
sb["ld"] = [(m.jnt_dofadr[mujoco_jid(env, "l_hip")],
m.jnt_dofadr[mujoco_jid(env, "l_knee")]),
(m.jnt_dofadr[mujoco_jid(env, "r_hip")],
m.jnt_dofadr[mujoco_jid(env, "r_knee")])]
# THE WAKE (the playground's first command): from the collapsed
# heap, run the baked rise - power on, lay out, tuck to launch,
# learned ascent, balance handover - then await orders
if (b.get("wake") and sb["fw"].mode != "wake"
and float(d.qpos[2]) < 0.30):
# wake arms whenever the robot is COLLAPSED - powered or
# not (the console keeps power:'on' every tick, so gating
# wake on the off state made the button dead: the heap was
# already powered by the time WAKE was pressed)
if sb.get("off"):
m.actuator_gainprm[:] = sb.pop("gain0")
m.actuator_biasprm[:] = sb.pop("bias0")
sb["off"] = False
sb["fw"].cmd_wake()
sb["leg"] = np.array([float(d.qpos[m.jnt_qposadr[
mujoco_jid(env, n)]])
for n in ("l_hip", "l_knee", "r_hip", "r_knee")])
# the power switch: idempotent gain kill / restore
pw = b.get("power")
if pw == "off" and not sb.get("off"):
sb["off"] = True
sb["gain0"] = m.actuator_gainprm.copy()
sb["bias0"] = m.actuator_biasprm.copy()
m.actuator_gainprm[:, 0] = 0.0
m.actuator_biasprm[:, 1:3] = 0.0
elif pw == "on" and sb.get("off"):
m.actuator_gainprm[:] = sb.pop("gain0")
m.actuator_biasprm[:] = sb.pop("bias0")
sb["off"] = False
# waking HOLDS the pose it wakes in: the leg command adopts
# the actual joints, and nothing moves until the operator
# moves something
sb["leg"] = np.array([float(d.qpos[m.jnt_qposadr[
mujoco_jid(env, n)]])
for n in ("l_hip", "l_knee", "r_hip", "r_knee")])
legs = [i for i, a in enumerate(env.spec["actuators"])
if a.get("mode") == "position"]
for i in legs:
env.act_center[i] = 0.5 * (env.ctrl_lo[i] + env.ctrl_hi[i])
env.act_span[i] = 0.5 * (env.ctrl_hi[i] - env.ctrl_lo[i])
if sb.get("leg") is None:
sb["leg"] = np.array([float(d.qpos[m.jnt_qposadr[
mujoco_jid(env, n)]])
for n in ("l_hip", "l_knee", "r_hip", "r_knee")])
stance_raw = b.get("stance_mm") # None = untouched = HOLD
want_v = float(np.clip(b.get("roam_v") or 0.0, -1.5, 1.5))
steer = float(np.clip(b.get("roam_w") or 0.0, -1.0, 1.0))
manual = b.get("legs_deg")
balance = b.get("balance", True)
# THE HAND RECORDER: while armed, every substep logs the commanded
# legs and toggle state - the operator performs a skill by hand,
# and the exam can replay the exact trajectory under identical
# physics. Demonstration beats description.
rec_cmd = b.get("record")
if rec_cmd == "start":
sb["rec"] = []
sb["rec_t0"] = float(env.t)
elif rec_cmd == "stop" and sb.get("rec") is not None:
out = ROOT / "runs" / "hand_demo.json"
json.dump({"frames": sb["rec"],
"note": "operator hand demonstration"},
open(out, "w"))
sb["rec"] = None
sb["rec_saved"] = len(json.load(open(out))["frames"])
hold_w = bool(b.get("hold_wheels"))
# the parking brake: on ENGAGE, remember exactly where the wheels
# are; while held, each wheel position-locks there - the base
# cannot roll away while poses are sculpted
if hold_w and sb.get("hold_q") is None:
sb["hold_q"] = [float(d.qpos[m.jnt_qposadr[mujoco_jid(env, n)]])
for n in ("l_wheel_j", "r_wheel_j")]
elif not hold_w:
sb["hold_q"] = None
push_n = float(b.get("push_n") or 0.0)
u = 0.0
nst = max(1, int(round(env.control_dt / m.opt.timestep)))
for _ in range(int(b.get("substeps", 3))):
d.xfrc_applied[sb["push_bid"], 0] = push_n
if sb.get("off"):
mujoco.mj_step(m, d, nstep=nst)
continue
# THE FIRMWARE (operator's architecture, now literal):
# free roam runs trainer/firmware.py - the same law objects
# the exams score, proved bit-exact at the refactor - and
# this server only plumbs sensors in and actuators out
hip_z = 0.5 * (float(d.xanchor[sb["hipj"][0]][2])
+ float(d.xanchor[sb["hipj"][1]][2]))
fw = sb["fw"]
if stance_raw is not None:
fw.cmd_stance(float(stance_raw))
fw.cmd_manual([math.radians(float(v)) for v in
(list(manual) + [0] * 4)[:4]]
if manual is not None else None)
# stick sign measured: positive roam_w yaws NEGATIVE
fw.stick(want_v, -steer * 1.2)
# sense: the same numbers, with the same meanings, as the
# exams read - the estimator, never ground truth
up = env.est_up
lean = math.atan2(-up[0], up[2])
rate = float(env.sensors.noisy(
"gyro", d.sensordata, env._raw)[1])
w = 0.0
for (wj, (hj, kj)) in zip(sb["wd"], sb["ld"]):
w += (float(d.qvel[wj]) + float(d.qvel[hj])
+ float(d.qvel[kj]) + rate)
speed = 0.5 * w * 0.0625
gz = float(env.sensors.noisy(
"gyro", d.sensordata, env._raw)[2])
rollr = math.atan2(up[1], max(1e-6, up[2]))
rrate = float(env.sensors.noisy(
"gyro", d.sensordata, env._raw)[0])
if "wb" not in sb:
sb["wb"] = [mujoco.mj_name2id(
m, mujoco.mjtObj.mjOBJ_BODY, n)
for n in ("l_wheel", "r_wheel")]
_xm = d.xmat[sb["push_bid"]].reshape(3, 3)
_fx, _fy = float(_xm[0][0]), float(_xm[1][0])
_fn = math.hypot(_fx, _fy) or 1.0
_cm = d.subtree_com[sb["push_bid"]]
_ap = 0.5 * (d.xpos[sb["wb"][0]] + d.xpos[sb["wb"][1]])
com_off = ((float(_cm[0]) - float(_ap[0])) * _fx
+ (float(_cm[1]) - float(_ap[1])) * _fy) / _fn
out = fw.tick({"lean": lean, "lrate": rate, "speed": speed,
"gz": gz, "rollr": rollr, "rrate": rrate,
"com_off": com_off, "hip_z": hip_z,
"cur_legs": sb["leg"]})
sb["mode"] = out["mode"]
sb["fwdbg"] = [round(sb["fw"].want_wz, 2),
round(sb["fw"].want_v, 2),
round(speed, 3),
round(getattr(sb["fw"], "hold_cmd", 0.0), 3),
(None if getattr(sb["fw"], "spin_rel", None)
is None else round(sb["fw"].spin_rel, 3)),
round(getattr(sb["fw"], "spin_settle", 0.0),
2),
round(gz, 3)]
if out["reseat"]:
# a failed wake re-seats to the heap for another try -
# with a DIFFERENT settle each time: the reseat is
# deterministic, so an identical heap replays the
# identical failure three times and gives up
LB.place_sitting(env)
LB.limp_settle(env, 0.6 + 0.25 * float(
getattr(sb["fw"].law, "tries", 1)))
sb["leg"] = np.array([float(d.qpos[m.jnt_qposadr[
mujoco_jid(env, n)]])
for n in ("l_hip", "l_knee", "r_hip", "r_knee")])
if out["leg_want"] is not None:
sb["leg"] = CTRL.slew(sb["leg"], out["leg_want"],
out["leg_rate"], env.control_dt)
u = float(out["u"])
t_t = float(out["t"])
# resting on the knee rollers = statically supported: no
# balancing wanted (the balance law dragged the powered heap
# 6.8m across the floor) - the wheels hold position until the
# rollers leave the ground, then balance takes over
# grounded = anything but the tyres touches the floor
# (measured: the heap rests on anonymous chassis and shank
# geoms, not the named rollers) - wheels-only contact is the
# balancing regime
grounded = False
for ci in range(d.ncon):
c = d.contact[ci]
if sb["floor_g"] == c.geom1:
other = c.geom2
elif sb["floor_g"] == c.geom2:
other = c.geom1
else:
continue
if other not in sb["tyre_g"]:
grounded = True
break
if grounded and sb.get("gnd_q") is None:
sb["gnd_q"] = [float(d.qpos[m.jnt_qposadr[
mujoco_jid(env, n)]])
for n in ("l_wheel_j", "r_wheel_j")]
elif not grounded:
sb["gnd_q"] = None
# wheel modes, in precedence order: HOLD (position-locked
# parking brake) > grounded auto-hold > balance > passive
if out["wake_drive"] is not None:
_drv = float(out["wake_drive"])
uw = [_drv, _drv]
u, t_t = 0.0, 0.0
sb["uw"] = uw
elif (not hold_w) and sb.get("gnd_q") is not None:
uw = []
for wi, n2 in enumerate(("l_wheel_j", "r_wheel_j")):
adr = m.jnt_qposadr[mujoco_jid(env, n2)]
dof = m.jnt_dofadr[mujoco_jid(env, n2)]
uw.append(float(np.clip(
-2.0 * (float(d.qpos[adr]) - sb["gnd_q"][wi])
- 0.2 * float(d.qvel[dof]), -1.0, 1.0)))
u, t_t = 0.0, 0.0
sb["uw"] = uw
elif hold_w and sb.get("hold_q") is not None:
uw = []
for wi, n2 in enumerate(("l_wheel_j", "r_wheel_j")):
adr = m.jnt_qposadr[mujoco_jid(env, n2)]
dof = m.jnt_dofadr[mujoco_jid(env, n2)]
uw.append(float(np.clip(
-3.0 * (float(d.qpos[adr]) - sb["hold_q"][wi])
- 0.3 * float(d.qvel[dof]), -1.0, 1.0)))
u, t_t = 0.0, 0.0
sb["uw"] = uw
elif balance:
# the firmware's wheel outputs pass through untouched
sb["uw"] = None
else:
sb["uw"] = None
u = 0.0
t_t = 0.0
act = np.zeros(env.act_dim)
for j, ai in enumerate(env.act_idx):
if ai in legs:
act[j] = np.clip(
(sb["leg"][legs.index(ai)] - env.act_center[ai])
/ max(1e-9, env.act_span[ai]), -1, 1)
else:
nm2 = env.spec["actuators"][ai]["name"]
if sb.get("uw") is not None:
act[j] = sb["uw"][0 if nm2.startswith("l") else 1]
else:
act[j] = float(np.clip(
u + (t_t if nm2.startswith("l") else -t_t),
-1, 1))
env.step(act)
if sb.get("rec") is not None:
sb["rec"].append({
"t": round(float(env.t) - sb["rec_t0"], 4),
"legs_deg": [round(math.degrees(float(v)), 2)
for v in sb["leg"]],
"hold": bool(hold_w), "balance": bool(balance),
"power": "off" if sb.get("off") else "on"})
# Report georeferenced position to Ghrups when driving in a target environment (origin
# supplied by the client). Throttled + threaded so it never stalls the sim tick.
_org = b.get("ghrups_origin")
if _org and _org.get("lat") is not None:
_mlat = 111320.0
_lat = _org["lat"] + float(env.base_pos[1]) / _mlat
_lng = _org["lng"] + float(env.base_pos[0]) / (_mlat * math.cos(math.radians(_org["lat"])))
_ghrups_report(_lat, _lng)
return self._send({
"frame": env.frame(),
"t": round(env.t, 3),
"height": round(env.height, 4),
"up_z": round(env.up_z, 4),
"pitch_deg": round(math.degrees(math.atan2(
-env.est_up[0], env.est_up[2])), 2),
"roll_deg": round(math.degrees(math.atan2(
env.est_up[1], max(1e-6, env.est_up[2]))), 2),
"tau": [round(float(v), 4) for v in d.actuator_force],
"qpos": [round(float(d.qpos[i]), 4)
for i in env.leg_joint_ids],
"fell": False,
"sandbox": True,
"power": "off" if sb.get("off") else "on",
"recording": len(sb["rec"]) if sb.get("rec") is not None else None,
"rec_saved": sb.pop("rec_saved", None),
"task": "bob",
"act_dim": env.act_dim,
"x": round(float(env.base_pos[0]), 4),
"y": round(float(env.base_pos[1]), 4),
"vx": round(float(env.vel_local[0]), 3),
"mode": sb.get("mode", "idle"),
"fwdbg": sb.get("fwdbg"),
"reqn": sb.setdefault("reqn", 0),
"ctl": {"u": round(u, 3),
"lean_deg": round(math.degrees(math.atan2(
-env.est_up[0], env.est_up[2])), 2)},
"wheel_rpm": [round(float(d.qvel[
m.jnt_dofadr[mujoco_jid(env, n)]]) * 60 / (2 * math.pi), 1)
for n in ("l_wheel_j", "r_wheel_j")],
"gyro": [round(float(v), 4) for v in env.sensors.noisy(
"gyro", d.sensordata, env._raw)],
"accel": [round(float(v), 3) for v in env.sensors.noisy(
"accel", d.sensordata, env._raw)],
})
def _live(self, b):
"""Step the hand-driven sim and return a frame.
Joint targets arrive in the same normalised [-1,1] the policy uses, so
what you feel on the sliders is exactly what the network commands.
"""
with LIVE["lock"]:
env = _live_env(reset=bool(b.get("reset")), task=b.get("task"),
tether=b.get("tether"), bank_deg=b.get("bank_deg"),
world=b.get("world"))
if b.get("record") is not None:
LIVE["rec"] = bool(b["record"])
if b.get("clear_demo"):
LIVE["obs"], LIVE["act"] = [], []
if b.get("save_demo"):
sys.path.insert(0, str(TRAINER))
import demo as demomod
path, total = demomod.save(b["save_demo"], LIVE["task"],
LIVE["obs"], LIVE["act"])
LIVE["obs"], LIVE["act"] = [], []
return self._send({"saved": str(path.name), "steps": total})
if b.get("reset"):
LIVE["odo"] = 0.0
LIVE["start_pose"] = b.get("pose")
import mjcf
ranges = []
for a in env.spec["actuators"]:
if a.get("mode") == "position":
j = mjcf.find_joint(env.spec, a["joint"])
ranges.append(j.get("range_deg", [-90, 90]))
else:
ranges.append(None)
# start in a named pose from the spec, e.g. sitting
pose = LIVE.get("start_pose")
poses = env.spec.get("poses", {})
if pose and pose in poses:
for jn, val in poses[pose].items():
if jn.startswith("note"):
continue
jid = mujoco_jid(env, jn)
if jid >= 0:
env.data.qpos[env.model.jnt_qposadr[jid]] = float(val)
aid = mujoco.mj_name2id(env.model,
mujoco.mjtObj.mjOBJ_ACTUATOR, jn + "_a")
if aid >= 0:
env.data.ctrl[aid] = float(val)
import mujoco as _mj
_mj.mj_forward(env.model, env.data)
for _ in range(400):
_mj.mj_step(env.model, env.data)
return self._send({"frame": env.frame(), "reset": True,
"ranges": ranges,
"props": env.props or [],
"poses": {k: v for k, v in poses.items()},
"nominal": {k: round(math.degrees(v), 1)
for k, v in env.spec["init"]["qpos"].items()},
"actuators": [a["name"] for a in env.spec["actuators"]],
"tau_limit": [float(x) for x in env.tau_limit],
"spec": {k: v for k, v in env.spec.items() if k != "_path"},
"bodies": env.body_names()})
if b.get("restand"):
# Put it back on its wheels, upright and still. Once it is over,
# the balance loop has nothing to work with - the wheels cannot
# right a robot lying on its side - so recovering has to be an
# explicit act, exactly like picking it up off the floor.
mm = float(b.get("stance_mm", 460.0)) / 1000.0
c = max(-1.0, min(1.0, (mm - 0.0625) / 0.400))
q = 2.0 * math.acos(c)
targets = {"l_hip": q / 2, "l_knee": -q, "r_hip": q / 2, "r_knee": -q}
env.data.qvel[:] = 0.0
env.data.qacc[:] = 0.0
for jn, val in targets.items():
jid = mujoco_jid(env, jn)
if jid >= 0:
env.data.qpos[env.model.jnt_qposadr[jid]] = val
aid = mujoco.mj_name2id(env.model, mujoco.mjtObj.mjOBJ_ACTUATOR,
jn + "_a")
if aid >= 0:
env.data.ctrl[aid] = val
if env.model.jnt_type[0] == mujoco.mjtJoint.mjJNT_FREE:
# back to the middle of the floor, not where it crashed:
# after a few runs it is metres away and off camera
env.data.qpos[0:3] = [0.0, 0.0, mm + 0.002]
env.data.qpos[3:7] = [1.0, 0.0, 0.0, 0.0] # perfectly upright
mujoco.mj_forward(env.model, env.data)
env._cache_truth(); env._est_update()
LIVE.pop("assist", None) # forget ramps and odometry
LIVE["obs_now"] = env._obs() if hasattr(env, "_obs") else LIVE.get("obs_now")
return self._send({"frame": env.frame(), "restand": True,
"height": round(env.height, 4),
"up_z": round(env.up_z, 4)})
if b.get("mode") == "gains":
# Show what a stage learned, the way that stage was scored.
# Replaying every stage through the balance law was useless:
# bob would never change height and recover would never be
# shoved, so you saw an identical standing robot each time.
stage = b.get("stage") or LEARN.get("stage", "balance")
# the learner's gains belong to the stage it was actually
# searching - replaying them through a DIFFERENT stage's
# law drove the jump demo with rotate's numbers (crouch
# clip pegged at 380, fired from a half-squat)
g = b.get("gains") or (LEARN.get("gains")
if LEARN.get("stage") == stage
else None)
if not g:
# round 0: the search has scored its seed but reported no
# round yet, which used to leave the watcher an inert
# straight-legged robot ("no gains yet", silently). The
# search STARTS from the stage seed - showing the seed law
# is showing the truth.
if str(TRAINER) not in sys.path:
sys.path.insert(0, str(TRAINER))
import learn_balance as _LBs
g = _LBs.STAGES.get(stage, {}).get("seed")
if not g and stage == "roam":
g = [0.0] # the benches carry roam entirely
if not g:
return self._send({"error": "no gains yet"}, 409)
g = list(np.array(g, dtype=float)) + [0.0] * 10
# A gradient, done by tilting gravity rather than the floor.
# For a body on a plane the two are equivalent, and this way
# nothing has to move: no respawning the robot on a ramp, no
# contact geometry to get wrong. Uphill is positive.
slope = math.radians(float(b.get("slope_deg", 0.0)))
env.model.opt.gravity[:] = [9.81 * math.sin(slope), 0.0,
-9.81 * math.cos(slope)]
if b.get("shove"):
LIVE.setdefault("gainstate", {})["pending_shove"] = float(b["shove"])
st = LIVE.setdefault("gainstate", {"odo": 0.0, "leg": None, "k": 0})
# THE UI AS AN INPUT TO THE LEARNING: a mouse-drag on the 3D
# robot arrives as a continuous force (push_n, Newtons).
# Each push is also RECORDED - peak force and duration - into
# runs/push_bank.json, and the bench exam replays the bank as
# its own case, so the search trains against the operator's
# actual hands. The bank is only written while no search is
# running, so a run's objective holds still mid-search.
# THE POWER SWITCH (operator's spec): off = every
# actuator's gain zeroed, the body goes truly limp and
# slumps wherever physics takes it; on = gains restored,
# the laws wake where it lies. Idempotent per state.
_pw = b.get("power")
if _pw == "off" and not st.get("power_off"):
st["power_off"] = True
st["pwr_gain"] = env.model.actuator_gainprm.copy()
st["pwr_bias"] = env.model.actuator_biasprm.copy()
env.model.actuator_gainprm[:, 0] = 0.0
env.model.actuator_biasprm[:, 1:3] = 0.0
env.min_h = 0.02 # a heap is not a fall
elif _pw == "on" and st.get("power_off"):
env.model.actuator_gainprm[:] = st.pop("pwr_gain")
env.model.actuator_biasprm[:] = st.pop("pwr_bias")
st["power_off"] = False
_pn = float(b.get("push_n") or 0.0)
st["push_n"] = _pn
_dtr = int(b.get("substeps", 3)) * env.control_dt
_rec = st.get("push_rec")
if _pn and not _rec:
st["push_rec"] = {"n": _pn, "dur": _dtr}
elif _pn and _rec:
_rec["dur"] += _dtr
if abs(_pn) > abs(_rec["n"]):
_rec["n"] = _pn
elif _rec and not _pn:
st.pop("push_rec", None)
if abs(_rec["n"]) >= 5.0 and _rec["dur"] >= 0.05:
if LEARN.get("running"):
# a push during a live search may not change the
# exam mid-run - it PENDS, and lands in the bank
# the moment the run ends, so it examines the
# NEXT run instead of destabilising this one
LIVE.setdefault("push_pending", []).append(_rec)
else:
_bank_pushes([_rec]
+ LIVE.pop("push_pending", []))
if b.get("zero") or st.get("stage") != stage:
if st.get("pwr_gain") is not None:
# leaving while powered off: restore the actuators
# or the model stays limp forever
env.model.actuator_gainprm[:] = st["pwr_gain"]
env.model.actuator_biasprm[:] = st["pwr_bias"]
# leaving slip/payload mid-demo must not leave the world
# slick or the chassis loaded
if st.get("mu0") is not None and len(st["mu0"]) == env.model.ngeom:
env.model.geom_friction[:, 0] = st["mu0"]
if st.get("m0") is not None and st.get("bid") is not None:
env.model.body_mass[st["bid"]] = st["m0"]
env.model.body_ipos[st["bid"]][0] = st["x0"]
st.clear()
st.update(odo=0.0, leg=None, k=0, stage=stage,
brace=None, v_cmd=0.0, u_hist=[])
# the base-layer reflex: recover's baked brace rides
# under every task stage, not just recover's own exam
try:
_rb = json.load(open(
ROOT / "runs" / "recover_gains.json"))["gains"]
st["reflex"] = (CTRL.ShoveBrace(_rb[8], _rb[9])
if len(_rb) >= 10
else CTRL.ShoveBrace(0.0, 0.3))
except (OSError, ValueError, IndexError):
st["reflex"] = CTRL.ShoveBrace(0.0, 0.3)
(kp_g, kd_g, stn_g, clamp_g, trim_s, leg_sp, shift_k,
shift_lim, brace_g, brace_f) = g[:10]
# Trim only where the SCORED law has it. balance's law has no
# trim term at all, but `trim_s or 1.0` used to turn its
# zero-padded slot into 1.0 - the demo held an offset the
# scorer never saw. stop zeroes it; level forces 1.0 below.
trim_s = ((trim_s or 1.0)
if stage in ("bob", "recover", "spin", "balance",
"balance_mid", "balance_crouch",
"stop", "roam", "rise")
else 0.0)
leg_sp = leg_sp or 1.0
if stage == "bob":
# bob's NEW 3-slot layout: trim scale and leg speed live
# in slots 0/1; the wheel gains come from the balance
# bakes, interpolated at the commanded height inside the
# substep loop below
trim_s, leg_sp = (g[0] or 1.0), (g[1] or 1.0)
if stage == "spin":
# spin's 5-slot layout: kinematic trim direct (slot 4 is
# the spin rate, NOT a trim multiplier), leg speed slot 0
trim_s, leg_sp = 1.0, (g[0] or 1.0)
legs = [i for i, a2 in enumerate(env.spec["actuators"])
if a2.get("mode") == "position"]
for i in legs:
env.act_center[i] = 0.5 * (env.ctrl_lo[i] + env.ctrl_hi[i])
env.act_span[i] = 0.5 * (env.ctrl_hi[i] - env.ctrl_lo[i])
jd = [env.model.jnt_dofadr[mujoco_jid(env, n)]
for n in ("l_wheel_j", "r_wheel_j")]
legjd = [(env.model.jnt_dofadr[mujoco_jid(env, "l_hip")],
env.model.jnt_dofadr[mujoco_jid(env, "l_knee")]),
(env.model.jnt_dofadr[mujoco_jid(env, "r_hip")],
env.model.jnt_dofadr[mujoco_jid(env, "r_knee")])]
if st["leg"] is None:
# Seed the ramp from where the legs ACTUALLY are - the
# same fix the hand-drive rig needed. Seeding at straight
# meant every restart yanked the legs straight before
# ramping back down, and early in a run the demo restarts
# every couple of seconds - so the screen showed an
# endless parade of straight-legged respawns.
st["leg"] = np.array([float(env.data.qpos[
env.model.jnt_qposadr[mujoco_jid(env, n)]])
for n in ("l_hip", "l_knee", "r_hip", "r_knee")])
rate_lim = CTRL.leg_rate(leg_sp)
GATES = [460, 400, 340]
if "push_bid" not in st:
st["push_bid"] = mujoco.mj_name2id(
env.model, mujoco.mjtObj.mjOBJ_BODY, "chassis")
for _ in range(int(b.get("substeps", 3))):
st["k"] += 1
# the operator's mouse force, live on the chassis
env.data.xfrc_applied[st["push_bid"], 0] = st.get(
"push_n", 0.0)
up = env.est_up
# full-quadrant: the old clamp saturated at +/-90 and
# blinded the far side of rotate's circle
lean = math.atan2(-up[0], up[2])
rate = float(env.sensors.noisy(
"gyro", env.data.sensordata, env._raw)[1])
# TRUE rolling speed: the encoder reads wheel-vs-shank,
# so leg motion must be added back (operator-diagnosed;
# same compensation as the scored law)
_w = 0.0
for (_wj, (_hj, _kj)) in zip(jd, legjd):
_w += (float(env.data.qvel[_wj])
+ float(env.data.qvel[_hj])
+ float(env.data.qvel[_kj]) + rate)
speed = 0.5 * _w * 0.0625
st["odo"] += speed * env.control_dt
# bob walks the height gates; bump crouches at its own
# searched stance; payload holds a working stance; the
# others hold 460
if b.get("stance_mm") is not None and stage != "spin":
stance = float(b["stance_mm"]) # you set the gate
# RELATIONAL follow, capped at the MEASURED survival
# ceiling: 350 deg/s square-wave yanks survive at
# 149mm drift, 600 deg/s falls in 4.5s with the
# current gain family. The yank case is in bob's
# exam now - when a bake proves faster, raise this.
rate_lim = max(rate_lim, math.radians(350.0))
elif stage == "bob":
stance = GATES[int(st["k"] / 120) % len(GATES)]
elif stage == "bump":
stance = float(np.clip(400.0 + 100.0 * g[6], 340.0, 462.0))
rate_lim = CTRL.leg_rate(g[7])
elif stage == "turn":
stance = float(np.clip(400.0 + 100.0 * g[12], 340.0, 462.0))
rate_lim = math.radians(120.0)
elif stage == "spin":
# gate changes and spinning are SEQUENCED, never
# simultaneous: a spin through the crouch descent
# is unbalanceable (the operator watched it degrade
# as it crouched). The gate is the OPERATOR'S
# slider now, like bob's; the stance slews first
# and the spin resumes when it lands.
_want_st = float(b.get("stance_mm") or 460.0)
_cur_st = st.get("sp_st", 460.0)
_cur_st += float(np.clip(_want_st - _cur_st,
-250.0 * env.control_dt,
250.0 * env.control_dt))
st["sp_st"] = _cur_st
stance = _cur_st
st["sp_hold"] = abs(_want_st - _cur_st) > 4.0
rate_lim = CTRL.leg_rate(g[0] or 1.0)
elif stage == "rotate":
stance = 420
rate_lim = CTRL.leg_rate(g[5] or 1.0)
elif stage == "stop":
# the brake crouch: stance dives toward the learned
# brake stance while braking, back up for the cruise
_ph = st.get("phase", 0)
_braking = (_ph % 320) >= 170
_bs = float(np.clip(
400.0 + 100.0 * (g[2] if len(g) > 2 else -0.2),
330.0, 462.0))
_tgt_mm = _bs if _braking else 430.0
_rmm = (300.0 * max(0.2, abs(g[3]
if len(g) > 3 else 1.0)) * env.control_dt)
st["stop_st"] = st.get("stop_st", 430.0) + float(
np.clip(_tgt_mm - st.get("stop_st", 430.0),
-_rmm, _rmm))
stance = st["stop_st"]
elif stage == "sit":
# the shutdown's own descending gate (one tick
# lagged - it is set in the want branch below)
stance = float(st.get("sit_st", 455.0))
elif stage == "jump":
# the jump's own gate (dive / absorb / recover)
stance = float(st.get("j_st", 455.0))
elif stage == "payload":
stance = 440
elif stage == "balance_mid":
stance = 400
elif stage == "balance_crouch":
stance = 340
else:
stance = 460
# recover slides the feet TOWARD the fall - the sign lives
# in controller.foot_shift; this file's copy had it
# inverted, so the demo fled every fall it was scored on
# catching
dx = (CTRL.foot_shift(shift_k, lean, abs(shift_lim or 0.1))
if stage == "recover" else 0.0)
if stage in ("spin", "bob"):
# about the measured CoM, same law as the exams
# (spin: slot 5; bob: slot 3)
if "com_ids" not in st:
st["com_ids"] = (
mujoco.mj_name2id(env.model,
mujoco.mjtObj.mjOBJ_BODY, "chassis"),
[mujoco.mj_name2id(env.model,
mujoco.mjtObj.mjOBJ_BODY, n)
for n in ("l_wheel", "r_wheel")])
_cid, _wid = st["com_ids"]
# offset projected on the robot's OWN fore-aft
# axis - world-x is wrong once the heading turns
_xm = env.data.xmat[_cid].reshape(3, 3)
_fx, _fy = float(_xm[0][0]), float(_xm[1][0])
_fn = math.hypot(_fx, _fy) or 1.0
_cm = env.data.subtree_com[_cid]
_ap = 0.5 * (env.data.xpos[_wid[0]]
+ env.data.xpos[_wid[1]])
_off = ((float(_cm[0]) - float(_ap[0])) * _fx
+ (float(_cm[1]) - float(_ap[1])) * _fy) / _fn
_kc = ((g[5] if len(g) > 5 else -1.2)
if stage == "spin"
else (g[3] if len(g) > 3 else -1.2))
st["fdx"] = st.get("fdx", 0.0) + (
env.control_dt / 0.08) * (
_kc * _off - st.get("fdx", 0.0))
dx = float(np.clip(st["fdx"], -0.14, 0.14))
# The level stage's whole point is DIFFERENT leg lengths
# - and the demo never applied them, so "use" on level
# showed a robot that could not level even in principle.
# It also gets the moving slope, as slewed gravity.
if stage == "level" and len(g) >= 8:
# level's slots differ from bob's: 4/5/6 are the
# levelling law, 7 is leg speed, and there is no trim
# multiplier at all. Reading them with bob's layout
# gave the demo a 3.5x trim bias and took its leg
# speed from the damping slot - a sluggish, lean-
# offset robot that was NOT the law being scored.
trim_s = 1.0
# same ramp law the scorer uses, INCLUDING the branch
# that opens to servo speed above |leg speed| 3 - the
# demo used to cap at 60 deg/s, so a fast-leg solution
# would have displayed slow
rate_lim = CTRL.leg_rate(g[7], floor=0.3,
unlimited_at=3.0)
# step-and-hold on the REAL table, the world the stage
# is scored on. The stale slew demo ran the new gains
# with the old sign convention: wrong leg up, leaning
# on flat - what you watched was a different robot.
_t = (st["k"] * env.control_dt) % 9.0
if _t < 1.0:
_fr = 0.0
elif _t < 2.0:
_fr = (_t - 1.0)
elif _t < 4.5:
_fr = 1.0
elif _t < 5.5:
_fr = 1.0 - 2.0 * (_t - 4.5)
else:
_fr = -1.0
env.data.ctrl[-1] = math.radians(6.0) * _fr
st["slope_deg"] = math.degrees(float(env.data.qpos[-1]))
# the BNO086: absolute tilt from the fused estimator,
# one reference for leveller, score and HUD
roll_n = math.atan2(env.est_up[1],
max(1e-6, env.est_up[2]))
st["spill_deg"] = abs(math.degrees(roll_n))
rr = float(env.sensors.noisy(
"gyro", env.data.sensordata, env._raw)[0])
dh = float(np.clip(g[4] * roll_n + g[5] * rr,
-abs(g[6]), abs(g[6])))
# the searched stance height and fore-aft foot shift,
# when the gains carry them (10-parameter level law)
if len(g) >= 10:
stance = float(np.clip(400.0 + 100.0 * g[8],
340.0, 462.0))
lean_n = math.atan2(-env.est_up[0],
max(1e-6, env.est_up[2]))
dx = float(np.clip(g[9] * lean_n, -0.08, 0.08))
hl, kl = leg_ik(max(0.20, min(0.462,
stance / 1000.0 + dh)), dx)
hr, kr = leg_ik(max(0.20, min(0.462,
stance / 1000.0 - dh)), dx)
want = np.array([hl, kl, hr, kr])
else:
env.model.opt.gravity[:] = [0.0, 0.0, -9.81]
st["slope_deg"] = 0.0
if stage == "turn" and len(g) >= 14:
# lean into the corner (inside leg squats - the
# measured sign): planned-turn feedforward, the
# leveller's roll loop, and the measured-force
# loop that drops the inside leg until the wheel
# loads rebalance.
# corner, settle, counter-corner, settle -
# the exit transient is the part worth watching
ph_t = (st["k"] * env.control_dt) % 14.0
yc = (0.7 if 3.0 <= ph_t < 6.0
else -0.7 if 8.0 <= ph_t < 11.0 else 0.0)
rr_t = float(env.sensors.noisy(
"gyro", env.data.sensordata, env._raw)[0])
roll_t = math.atan2(env.est_up[1],
max(1e-6, env.est_up[2]))
# low-passed lateral-accel command (tau 0.25s):
# the lean walks in with the corner and recovers
# gradually after it, same as the scored law
# curvature governor: the learned lateral
# budget widens too-tight commands to the
# holdable radius for the current speed
_vf = st.get("v_f", 0.0)
yce = yc
if abs(_vf) > 0.15 and abs(yc) > 1e-6:
_wcap = max(0.1, abs(g[13])) / abs(_vf)
yce = float(np.clip(yc, -_wcap, _wcap))
st["yce"] = yce
st["ff_lat"] = st.get("ff_lat", 0.0) + (
env.control_dt / 0.25) * (
_vf * yce - st.get("ff_lat", 0.0))
if "touch_adr" not in st:
st["touch_adr"] = [env.model.sensor_adr[
mujoco.mj_name2id(env.model,
mujoco.mjtObj.mjOBJ_SENSOR, nm)]
for nm in ("l_touch", "r_touch")]
_fl = float(env.data.sensordata[st["touch_adr"][0]])
_fr = float(env.data.sensordata[st["touch_adr"][1]])
_imb = (_fr - _fl) / max(25.0, _fr + _fl)
# lean machinery gated on the turn command, as
# scored: standing means still legs
st["gate"] = st.get("gate", 0.0) + (
env.control_dt / 0.3) * (
(1.0 if abs(yc) > 0.05 else 0.0)
- st.get("gate", 0.0))
# feedback TRIMS through a 0.4s low-pass (the
# sweet-spot settle); the feedforward is fast
_fb = (g[9] * roll_t + g[10] * rr_t
+ g[11] * _imb)
st["fb_slow"] = st.get("fb_slow", 0.0) + (
env.control_dt / 0.4) * (
_fb - st.get("fb_slow", 0.0))
dh_t = float(np.clip(
st["gate"] * (g[8] * st["ff_lat"]
+ st["fb_slow"]),
-0.05, 0.05))
hl2, kl2 = leg_ik(max(0.20, min(0.462,
stance / 1000.0 + dh_t)))
hr2, kr2 = leg_ik(max(0.20, min(0.462,
stance / 1000.0 - dh_t)))
want = np.array([hl2, kl2, hr2, kr2])
elif stage == "rise" and len(g) >= 4:
# v41: choreography to the LAUNCH pose, then
# the LEARNED ascent (drive+taper+lift) with
# the discovered balance handover height
if b.get("reseat"):
st["seated"] = False
if not st.get("seated"):
if str(TRAINER) not in sys.path:
sys.path.insert(0, str(TRAINER))
import learn_balance as _LBr
env.min_h = 0.02
_LBr.place_sitting(env)
_LBr.limp_settle(env, 1.0)
st["seated"] = True
st["rph"] = 0
st["leg"] = np.array([
float(env.data.qpos[
env.model.jnt_qposadr[
mujoco_jid(env, n)]])
for n in ("l_hip", "l_knee",
"r_hip", "r_knee")])
st["rlock"] = [float(env.data.qpos[
env.model.jnt_qposadr[
mujoco_jid(env, n)]])
for n in ("l_wheel_j", "r_wheel_j")]
import learn_balance as _LBr2
_lay = np.array([
math.radians(_LBr2.RISE_LAY[0]),
math.radians(_LBr2.RISE_LAY[1])] * 2)
_launch = np.array([
math.radians(_LBr2.RISE_LAUNCH[0]),
math.radians(_LBr2.RISE_LAUNCH[1])] * 2)
_sth, _stk = CTRL.stance_to_leg(0.455)
_stand = np.array([_sth, _stk] * 2)
_slow = math.radians(35.0)
_liftr = math.radians(
150.0 * min(1.5, max(0.1, abs(g[2]))))
_hj3 = [mujoco_jid(env, n)
for n in ("l_hip", "r_hip")]
_hz3 = 0.5 * (
float(env.data.xanchor[_hj3[0]][2])
+ float(env.data.xanchor[_hj3[1]][2]))
ph = st.get("rph", 0)
if ph == 0:
want, rate_lim = _lay, _slow
st["rise_hold"] = True
st["rise_drive"] = None
if (st["leg"] is not None and float(
np.max(np.abs(st["leg"] - _lay)))
< math.radians(5.0)):
st["rph"] = 1
elif ph == 1:
want, rate_lim = _launch, _slow
st["rise_hold"] = True
st["rise_drive"] = None
if float(np.max(np.abs(
st["leg"] - _launch))) < math.radians(5.0):
st["rph"] = 2
else:
want, rate_lim = _stand, _liftr
st["rise_hold"] = False
if 1000.0 * _hz3 < max(
200.0, min(462.0, abs(g[3]))):
st["rise_drive"] = float(np.clip(
g[0] + g[1]
* (_hz3 - 0.165) / 0.29,
-1.0, 1.0))
else:
st["rise_drive"] = None
elif stage == "sit" and len(g) >= 2:
# the SHUTDOWN, v53: the wheels start on the
# other side to normal running, so it's the
# gate in reverse - bob's law verbatim in the
# MIRROR knee fold rides the height gate all
# the way down, then the wheels lock, the
# legs fold to the heap and the POWER CUTS.
# Stays down until reseat - off is the point.
if b.get("reseat"):
st["sph"] = None
if st.get("sph") is None:
if st.get("pwr_gain") is not None:
env.model.actuator_gainprm[:] = (
st["pwr_gain"])
env.model.actuator_biasprm[:] = (
st["pwr_bias"])
st["pwr_gain"] = None
st["pwr_bias"] = None
st["sph"] = 0
st["sit_st"] = 455.0
st["sit_fdx"] = 0.0
env.reset()
_sh, _sk = CTRL.stance_to_leg(
0.455, 0.0, branch=+1.0)
for _nm, _v in (("l_hip", _sh),
("l_knee", _sk),
("r_hip", _sh),
("r_knee", _sk)):
env.data.qpos[env.model.jnt_qposadr[
mujoco_jid(env, _nm)]] = _v
env.data.qpos[0:3] = [0.0, 0.0, 0.4575]
env.data.qpos[3:7] = [1, 0, 0, 0]
env.data.qvel[:] = 0
import mujoco as _mj
_mj.mj_forward(env.model, env.data)
env.min_h = 0.02
st["leg"] = np.array([
_sh, _sk, _sh, _sk])
if str(TRAINER) not in sys.path:
sys.path.insert(0, str(TRAINER))
import learn_balance as _LBs2
_bobg = (list(_LBs2.accepted("bob")
or [1.0, 1.0, 0.0, -0.5])
+ [1.0, 1.0, 0.0, -0.5])[:4]
st["sit_trims"] = _bobg[0] or 1.0
_heap = np.array([
_LBs2.SIT_POSE["l_hip"],
_LBs2.SIT_POSE["l_knee"]] * 2)
sph = st["sph"]
st["rise_hold"] = False
st["rise_drive"] = None
LIVE["dbg"] = {"sph": sph,
"sit_st": round(
st.get("sit_st", -1), 1),
"off": st.get("pwr_gain")
is not None}
if sph == 0:
st["sit_st"] = max(170.0, st.get(
"sit_st", 455.0)
- 300.0 * min(1.0, max(0.08,
abs(g[0])))
* env.control_dt)
# bob's realtime CoM hold, bob's slot 3
_cid2 = mujoco.mj_name2id(
env.model, mujoco.mjtObj.mjOBJ_BODY,
"chassis")
_wid2 = [mujoco.mj_name2id(
env.model, mujoco.mjtObj.mjOBJ_BODY,
n) for n in ("l_wheel", "r_wheel")]
_com2 = float(
env.data.subtree_com[_cid2][0])
_ax2 = 0.5 * (
float(env.data.xpos[_wid2[0]][0])
+ float(env.data.xpos[_wid2[1]][0]))
st["sit_fdx"] = st.get("sit_fdx", 0.0) + (
env.control_dt / 0.04) * (
_bobg[3] * (_com2 - _ax2)
- st.get("sit_fdx", 0.0))
_dxb = float(np.clip(
_bobg[2] * (0.46
- st["sit_st"] / 1000.0)
+ st["sit_fdx"], -0.06, 0.14))
_dh2, _dk2 = CTRL.stance_to_leg(
st["sit_st"] / 1000.0, _dxb,
branch=+1.0)
want = np.array([_dh2, _dk2] * 2)
rate_lim = CTRL.leg_rate(_bobg[1])
if st["sit_st"] <= 171.0:
st["sph"] = 1
st["rlock"] = [float(
env.data.qpos[
env.model.jnt_qposadr[
mujoco_jid(env, n2)]])
for n2 in ("l_wheel_j",
"r_wheel_j")]
else:
want = _heap
rate_lim = math.radians(
150.0 * min(1.0, max(0.08,
abs(g[1]))))
st["rise_hold"] = True
if (st.get("pwr_gain") is None
and st["leg"] is not None
and float(np.max(np.abs(
st["leg"] - _heap)))
< math.radians(10.0)):
# THE POWER CUTS - limp in its heap,
# exactly as the real one ends
st["pwr_gain"] = (
env.model
.actuator_gainprm.copy())
st["pwr_bias"] = (
env.model
.actuator_biasprm.copy())
env.model.actuator_gainprm[:, 0] = 0.0
env.model.actuator_biasprm[
:, 1:3] = 0.0
elif stage == "jump":
# v56: up only, launched along the measured
# tether trajectory (JUMP_TRAJ) - crouch,
# level, fire up the path, tuck at apex,
# absorb. One button.
if str(TRAINER) not in sys.path:
sys.path.insert(0, str(TRAINER))
import learn_balance as _LBj
_bobj = (list(_LBj.accepted("bob") or [1.0])
+ [1.0])[:1]
st["j_trims"] = _bobj[0] or 1.0
_cr2 = float(np.clip(300.0 + 80.0 * g[0],
218.0, 380.0))
_tk2 = float(np.clip(340.0 + 80.0 * g[3],
250.0, 430.0))
_ab2 = float(np.clip(300.0 + 80.0 * g[4],
230.0, 400.0))
_dvr = 300.0 * min(1.0, max(0.08, abs(g[1])))
_lnr = 2500.0 * min(1.0, max(0.1, abs(g[2])))
_rcr = 300.0 * min(1.0, max(0.08, abs(g[5])))
# v94/v95 launch slots (zero = the old
# hardcoded constants, so 19-slot bakes fly
# exactly as before). This block used to BE
# the second implementation the NEXT notes
# warn about: the exam learned a slower
# fire and a gyro-fed wheel brake, and the
# demo flew the old constants - backflip.
_g24 = (lambda i: float(g[i])
if len(g) > i else 0.0)
_jfr = 2000.0 * (1.0 + float(np.clip(
_g24(19), -0.7, 1.0)))
_jtop = float(np.clip(
458.0 + 40.0 * _g24(20), 380.0, 461.0))
_jgate = 0.15 * (1.0 + float(np.clip(
_g24(21), -0.9, 2.0)))
_jbrace = 0.1 * (1.0 + float(np.clip(
_g24(22), -0.8, 3.0)))
_jsteer = _g24(23)
# v96 slot 25: the hip tracks the shin's
# fold rate through tuck and brace (the
# operator's cancellation)
_jtbal = _g24(24)
_jkdof = env.model.jnt_dofadr[
mujoco_jid(env, "l_knee")]
if (b.get("jump_dir")
and st.get("jph") is None):
st["jph"] = 0
st["jph_t"] = 0.0
st["j_free"] = 0.0
st["j_am"] = 0
st["j_land_k"] = None
jph = st.get("jph")
LIVE["dbg"] = {"jph": jph,
"j_st": round(st.get(
"j_st", -1.0), 1)}
st["j_trimf"] = 1.0
# the fire lean belongs to the launch beat
# ONLY - it used to persist after a jump
# ended and the idle balance stood leaning
# 28 deg, driving away (operator: 'it isn't
# balancing to begin with')
st["j_flean"] = 0.0
st["uw_pair"] = None
st["j_catch_u"] = 0.0
if not st.get("j_seated"):
# a clean LEVEL entry: the view used to
# inherit whatever pose the last stage
# left (one wheel cocked off the floor)
st["j_seated"] = True
env.reset()
_sh5, _sk5 = CTRL.stance_to_leg(0.455)
for _nm5, _v5 in (("l_hip", _sh5),
("l_knee", _sk5),
("r_hip", _sh5),
("r_knee", _sk5)):
env.data.qpos[
env.model.jnt_qposadr[
mujoco_jid(env, _nm5)]] = _v5
env.data.qpos[0:3] = [0.0, 0.0, 0.4575]
env.data.qpos[3:7] = [1, 0, 0, 0]
env.data.qvel[:] = 0
import mujoco as _mj5
_mj5.mj_forward(env.model, env.data)
st["leg"] = np.array(
[_sh5, _sk5, _sh5, _sk5])
if "j_floor" not in st:
st["j_floor"] = mujoco.mj_name2id(
env.model, mujoco.mjtObj.mjOBJ_GEOM,
"floor")
if "j_wids" not in st:
st["j_wids"] = [mujoco.mj_name2id(
env.model, mujoco.mjtObj.mjOBJ_BODY,
n) for n in ("l_wheel", "r_wheel")]
_jcon = False
for _ci in range(env.data.ncon):
_cc = env.data.contact[_ci]
if st["j_floor"] in (_cc.geom1,
_cc.geom2):
_jcon = True
break
_wq2 = [env.model.jnt_qposadr[
mujoco_jid(env, n2)]
for n2 in ("l_wheel_j", "r_wheel_j")]
_wd2 = [env.model.jnt_dofadr[
mujoco_jid(env, n2)]
for n2 in ("l_wheel_j", "r_wheel_j")]
want = None
rate_lim = math.radians(400.0)
if jph == 0:
st["j_st"] = max(_cr2, st.get(
"j_st", 455.0)
- _dvr * env.control_dt)
if st["j_st"] <= _cr2 + 0.5:
st["jph"] = 1
st["jph_t"] = 0.0
elif jph == 1:
st["jph_t"] += env.control_dt
st["j_dxs"] = float(np.clip(
0.05 * g[11], -0.06, 0.06))
# ramp INTO the fire lean (slot 9) -
# the deliberate forward lean that
# spends the rotation budget before
# the launch can backflip
st["j_trimf"] = max(
0.0, 1.0 - st["jph_t"] / 0.15)
st["j_flean"] = (math.radians(
6.0 * g[9])
* (1.0 - st["j_trimf"]))
_fl = math.radians(6.0 * g[9])
# v101: planted feet before firing -
# exam law verbatim (drift at fire
# comes back down as the roll-out)
if (st["jph_t"] >= 0.1
and abs(float(
env.data.qvel[2])) < _jgate
and (abs(speed) < 0.1
or st["jph_t"] > 2.5)
and (abs(lean - _fl)
< math.radians(5.0)
or st["jph_t"] > 2.0)):
st["jph"] = 2
st["jph_t"] = 0.0
st["j_hcmd"] = _cr2
st["j_lock"] = [float(
env.data.qpos[a]) for a in _wq2]
elif jph == 2:
# the fire: wheels DRIVE forward while
# the legs open along the tether path
# at full rate (operator's spec) - the
# feet chase the CoM, no backward pin
st["jph_t"] += env.control_dt
st["j_hcmd"] = min(_jtop, st["j_hcmd"]
+ _jfr
* env.control_dt)
# the hip ratio (slot 11): knee drives
# the path, hip releases from its TRUE
# crouch angle at a fraction of the
# knee's pace
_ckd = float(np.clip(
0.1 * g[15], -0.15, 0.15))
_jk3 = CTRL.stance_to_leg(
st["j_hcmd"] / 1000.0,
_ckd * st["j_hcmd"] / 1000.0)[1]
_ch2 = CTRL.stance_to_leg(
_cr2 / 1000.0,
_ckd * _cr2 / 1000.0)[0]
_eh2 = CTRL.stance_to_leg(
0.462, _ckd * 0.462)[0]
_ck2 = CTRL.stance_to_leg(
_cr2 / 1000.0,
_ckd * _cr2 / 1000.0)[1]
_ek2 = CTRL.stance_to_leg(
0.462, _ckd * 0.462)[1]
_ka2 = float(env.data.qpos[
env.model.jnt_qposadr[
mujoco_jid(env, "l_knee")]])
_pg2 = float(np.clip(
(_ka2 - _ck2) / (_ek2 - _ck2),
0.0, 1.0))
_hl2 = float(np.clip(g[12], 0.0, 0.95))
_rl2 = max(0.0, (_pg2 - _hl2)
/ max(0.05, 1.0 - _hl2))
_jh3 = _ch2 + min(1.0, (g[10] / 0.45)
* _rl2) * (
_eh2 - _ch2)
_jk3 = _jk3 + float(np.clip(
g[13] * (_jk3 - _ka2), -1.2, 1.2))
_jh3 += math.radians(8.0 * g[7])
_jh3 += float(np.clip(
g[8] * rate, -0.35, 0.35))
want = np.array([_jh3, _jk3] * 2)
rate_lim = math.radians(3000.0)
# fire steer: the drive answers the
# felt pitch rate mid-thrust, same as
# the exam law
st["uw_pair"] = [float(np.clip(
3.0 * g[2] + _jsteer * rate,
-1, 1))] * 2
if not _jcon and st["jph_t"] > 0.02:
st["jph"] = 3
st["j_air"] = 0.0
st["j_free"] = 0.0
st["j_am"] = 1
elif jph == 3:
# three beats, the scored law verbatim:
# stretch -> tuck at the searched knee
# moment (the counter-rotation pulse)
# -> brace for impact
st["j_air"] = (st.get("j_air", 0.0)
+ env.control_dt
if not _jcon else 0.0)
st["j_free"] = max(st.get("j_free", 0.0),
st["j_air"])
_am = st.get("j_am", 0)
_ka4 = float(env.data.qpos[
env.model.jnt_qposadr[
mujoco_jid(env, "l_knee")]])
_ek4 = CTRL.stance_to_leg(0.462)[1]
_tr4 = math.radians(1000.0 * min(
2.0, max(0.3, abs(g[14]))))
if _am == 0:
_jh3, _jk3 = CTRL.stance_to_leg(
0.462)
rate_lim = math.radians(900.0)
_ta4 = _ek4 - math.radians(
15.0 - 15.0 * g[16])
if (_ka4 > _ta4
or float(env.data
.qvel[2]) < 0.0):
st["j_am"] = 1
elif _am == 1:
_jh3, _jk3 = CTRL.stance_to_leg(
_tk2 / 1000.0)
_jh3 += g[17]
_jh3 += float(np.clip(
_jtbal * float(env.data.qvel[
_jkdof]), -0.7, 0.7))
_jk3 += float(np.clip(
g[18] * rate, -0.6, 0.6))
rate_lim = math.radians(3000.0)
_vz4 = float(env.data.qvel[2])
_hw4 = min(
float(env.data.xpos[
st["j_wids"][0]][2]),
float(env.data.xpos[
st["j_wids"][1]][2])) - 0.0625
if (_vz4 < -0.05
and _hw4 / max(0.05, -_vz4)
< _jbrace):
st["j_am"] = 2
else:
_jh3, _jk3 = CTRL.stance_to_leg(
_ab2 / 1000.0)
_jh3 += float(np.clip(
_jtbal * float(env.data.qvel[
_jkdof]), -0.7, 0.7))
_jk3 += float(np.clip(
g[18] * rate, -0.6, 0.6))
rate_lim = _tr4
want = np.array([_jh3, _jk3] * 2)
# land lean (slot 32): flight attitude
# aims at a forward lean, exam verbatim
_att = float(np.clip(
g[6] * ((lean - math.radians(
6.0 * _g24(31))) * 8.0 + rate),
-1.0, 1.0))
# air spin (slot 27): reaction-wheel
# torque tied to the shin fold rate,
# exam law verbatim - the position
# lock fades as the spin engages
_wsp = float(np.clip(
_g24(26) * float(env.data.qvel[
_jkdof]), -1.0, 1.0))
_wlk = max(0.0, 1.0
- min(1.0, abs(_g24(26))))
st["uw_pair"] = [float(np.clip(
_wlk * (-1.5 * (float(env.data.qpos[
_wq2[wi]]) - st["j_lock"][wi])
- 0.15 * float(env.data.qvel[
_wd2[wi]]))
+ _wsp + _att, -1, 1))
for wi in (0, 1)]
if _jcon and st.get("j_free", 0.0) > 0.06:
st["jph"] = 4
st["j_st"] = _ab2
st["j_land_k"] = st["k"]
elif jph == 4:
st["j_st"] = min(455.0, st.get(
"j_st", _ab2)
+ _rcr * env.control_dt)
if st["j_st"] >= 454.5:
st["jph"] = None
else:
st["j_st"] = min(455.0, st.get(
"j_st", 455.0)
+ _rcr * env.control_dt)
# the landing reflexes run on their own 3s
# clock from touchdown, OUTSIDE the phase
# chain (jph 4 lasts a fraction of a second
# - gating on it cut the catch short):
# land catch (slot 26, 0.4s fade) + roll
# fix (slot 28, measured ground speed, 3s
# fade), exam law verbatim
st["j_catch_u"] = 0.0
if st.get("j_land_k") is not None:
_lel = ((st["k"] - st["j_land_k"])
* env.control_dt)
if _lel < 3.0:
st["j_catch_u"] = (
float(np.clip(
_g24(25) * (8.0 * lean
+ rate),
-1.0, 1.0))
* max(0.0, 1.0 - _lel / 0.4)
+ float(np.clip(
_g24(27) * speed,
-1.0, 1.0))
* max(0.0, 1.0 - _lel / 3.0))
if want is None:
_jh2, _jk2 = CTRL.stance_to_leg(
st.get("j_st", 455.0) / 1000.0)
# the catch's legs (slots 29-31),
# exam law verbatim: hip offset + hip
# and knee answer the felt lean over
# the 3s landing window
if st.get("j_land_k") is not None:
_le2 = ((st["k"] - st["j_land_k"])
* env.control_dt)
_cf2 = max(0.0, 1.0 - _le2 / 3.0)
_jh2 += float(np.clip(
0.5 * _g24(28)
+ _g24(29) * lean,
-0.5, 0.5)) * _cf2
_jk2 += float(np.clip(
_g24(30) * lean,
-0.5, 0.5)) * _cf2
want = np.array([_jh2, _jk2] * 2)
elif stage == "rotate" and len(g) >= 6:
# full-circle chassis rotation about the hip
# axis: the hip IS the gaze joint. Slider sends
# rot_deg (+/-360, positive = nose down); the
# command chases it at the law's own rate,
# slowed through the +/-90 lever arcs.
if b.get("rot_deg") is not None:
tgt = math.radians(float(b["rot_deg"]))
else:
ph_l = (st["k"] * env.control_dt) % 24.0
tgt = (math.radians(90.0) if 3 <= ph_l < 8
else math.radians(180.0)
if 8 <= ph_l < 14
else math.radians(-90.0)
if 16 <= ph_l < 21 else 0.0)
_rc = st.get("rot_cmd", 0.0)
rr = max(0.05, abs(g[0])) * (
1.0 - min(0.75, abs(g[1]))
* abs(math.sin(_rc)))
st["rot_cmd"] = _rc + float(
np.clip(tgt - _rc, -rr * env.control_dt,
rr * env.control_dt))
# realtime CoM hold: legs walk the feet to
# keep the axle under the measured total CoM
if "com_ids" not in st:
st["com_ids"] = (
mujoco.mj_name2id(env.model,
mujoco.mjtObj.mjOBJ_BODY, "chassis"),
[mujoco.mj_name2id(env.model,
mujoco.mjtObj.mjOBJ_BODY, n)
for n in ("l_wheel", "r_wheel")])
_cid, _wid = st["com_ids"]
_com = float(env.data.subtree_com[_cid][0])
_ax = 0.5 * (float(env.data.xpos[_wid[0]][0])
+ float(env.data.xpos[_wid[1]][0]))
st["fdx"] = st.get("fdx", 0.0) + (
env.control_dt / 0.08) * (
g[2] * (_com - _ax) - st.get("fdx", 0.0))
dxl = float(np.clip(
g[3] + st["fdx"]
- float(np.clip(abs(g[4]) * st["odo"],
-0.08, 0.08)),
-0.14, 0.14))
# the mirror rule: face flipped past +/-90 =
# knee folds the other way, staying toward the
# back of the new facing (hysteresis at 95/85)
_br = st.get("branch", -1.0)
_c = math.cos(st["rot_cmd"])
if _br < 0 and _c < -0.087:
_br = +1.0
elif _br > 0 and _c > 0.087:
_br = -1.0
st["branch"] = _br
# straight-leg crossing (operator's spec): the
# leg extends to full height as the rotation
# CROSSES +/-90, where the two knee folds meet
# at the straight-leg singularity - continuous
# swap. Only when a crossing lies AHEAD; a
# boundary at the target is a destination.
_mm2 = math.pi / 2.0
_rc2 = st["rot_cmd"]
_lo2, _hi2 = sorted(
(_rc2 - 0.3 * (1 if tgt >= _rc2 else -1),
tgt))
_xa = any(
nn % 2 != 0 and abs(nn * _mm2 - tgt) > 0.12
for nn in range(
math.ceil(_lo2 / _mm2 - 1e-9),
math.floor(_hi2 / _mm2 + 1e-9) + 1))
_bl2 = (abs(math.sin(_rc2)) ** 4) if _xa else 0.0
_steff = stance + (462.0 - stance) * _bl2
# legs hold the stance; the hip carries the
# rotation (nose-down theta = hip0 - theta,
# measured 13 Aug)
_h0, _k0 = CTRL.stance_to_leg(
_steff / 1000.0, dxl, _br)
want = np.array([_h0 - st["rot_cmd"], _k0,
_h0 - st["rot_cmd"], _k0])
rate_lim = CTRL.leg_rate(g[5] or 1.0)
elif stage == "roam" and b.get("legs_deg") is not None:
# FULL MANUAL LIMBS (the operator's lab): the
# cockpit commands each joint directly while
# the wheels keep balancing - experiment by
# hand, then teach what worked
_ld = [math.radians(float(v))
for v in (list(b["legs_deg"]) + [0] * 4)[:4]]
want = np.array(_ld)
rate_lim = math.radians(240.0)
else:
hp, kn = leg_ik(stance / 1000.0, dx)
want = np.array([hp, kn, hp, kn])
if stage == "spin" and len(g) >= 8:
# THE BANK (slots 7/8 feedback + slot 10
# feedforward): roll answered through the legs,
# and the held part flips with spin DIRECTION
# (gyroscopic torque flips with yaw) - the
# feedback-only demo showed the wrong leg down
# in anticlockwise spins (operator caught it)
_rollb = math.atan2(env.est_up[1],
max(1e-6, env.est_up[2]))
_gyb = env.sensors.noisy(
"gyro", env.data.sensordata, env._raw)
_rrb = float(_gyb[0])
_gzb = float(_gyb[2])
_dhb = float(np.clip(
g[6] * _rollb + g[7] * _rrb
+ (g[9] if len(g) > 9 else 0.0) * _gzb,
-0.05, 0.05))
# the bank sleeps at standstill - fades in
# with actual spin rate
_dhb *= min(1.0, abs(_gzb) / 0.3)
# the skater (slot 11): height rises with rate
_skb = 100.0 * abs(g[10] if len(g) > 10
else 0.0) * abs(_gzb)
_dxb2 = float(np.clip(st.get("fdx", 0.0),
-0.14, 0.14))
_stk = min(0.460, (stance + _skb) / 1000.0)
_hlb, _klb = leg_ik(max(0.2, min(
0.462, _stk + _dhb)), _dxb2)
_hrb, _krb = leg_ik(max(0.2, min(
0.462, _stk - _dhb)), _dxb2)
want = np.array([_hlb, _klb, _hrb, _krb])
if stage in ("roam", "rise"):
_hj2 = [mujoco_jid(env, n)
for n in ("l_hip", "r_hip")]
stance = float(np.clip(
500.0 * (float(env.data.xanchor[_hj2[0]][2])
+ float(env.data.xanchor[_hj2[1]][2])),
340.0, 462.0))
if stage in ("bob", "rotate", "stop", "roam",
"spin", "rise", "sit", "jump"):
# the benches' proven wheel gains at the stance the
# law is commanding RIGHT NOW (look recomputes its
# stance from the gaze above, so this sits after) -
# same interpolation the scored laws use
if "bpts" not in st:
if str(TRAINER) not in sys.path:
sys.path.insert(0, str(TRAINER))
import learn_balance as _LB
st["bpts"] = _LB.balance_points()
if st["bpts"]:
import learn_balance as _LB
kp_g, kd_g, stn_g, clamp_g = [
float(v) for v in _LB.interp_wheel_gains(
st["bpts"], stance)]
else:
# no bench baked yet: show the bench seed rather
# than a limp robot - the EXAM still refuses to
# run, so nothing unproven can be baked from this
kp_g, kd_g, stn_g, clamp_g = 3.65, 2.91, 1.02, 0.52
st["leg"] = CTRL.slew(st["leg"], want, rate_lim,
env.control_dt)
trim = math.radians(lean_trim_deg(stance)) * trim_s
if stage == "sit":
# the mirror fold balances tipped the OTHER way:
# its own measured trim curve, bob's trim scale
import learn_balance as _LBt
trim = (math.radians(
_LBt.mirror_trim_deg(stance))
* st.get("sit_trims", 1.0))
if stage == "spin" and len(g) > 8:
# cent trim (slot 9): the equilibrium lean
# shifts with yaw-rate squared - same term the
# exam scores
_gzc = float(env.sensors.noisy(
"gyro", env.data.sensordata, env._raw)[2])
trim += float(np.clip(g[8] * _gzc * _gzc,
-0.25, 0.25))
if stage == "jump":
# normal fold, bob's trim scale; the pre-fire
# beat blends from balance trim to the fire
# lean (slot 9's deliberate forward lean)
trim = (math.radians(lean_trim_deg(stance))
* st.get("j_trims", 1.0)
* st.get("j_trimf", 1.0)
+ st.get("j_flean", 0.0))
if stage in ("bump", "turn"):
trim = math.radians(lean_trim_deg(stance))
if stage == "rotate":
# the rotation IS the pitch target - no legacy trim
trim = -st.get("rot_cmd", 0.0)
if stage == "payload":
# the load appears and disappears on a cycle; the
# searched integral re-finds the trim from the drift
bid = mujoco.mj_name2id(env.model,
mujoco.mjtObj.mjOBJ_BODY,
"chassis")
if st.get("m0") is None:
st["bid"] = bid
st["m0"] = float(env.model.body_mass[bid])
st["x0"] = float(env.model.body_ipos[bid][0])
ph = (st["k"] * env.control_dt) % 10.0
kg = 2.0 if ph >= 2.0 else 0.0
env.model.body_mass[bid] = st["m0"] + kg
env.model.body_ipos[bid][0] = (
(st["m0"] * st["x0"] + kg * 0.04)
/ (st["m0"] + kg))
if st.get("z0") is None:
st["z0"] = float(
env.model.body_ipos[bid][2])
env.model.body_ipos[bid][2] = (
(st["m0"] * st["z0"] + kg * 0.13)
/ (st["m0"] + kg))
st["loaded_kg"] = kg
ki = abs(g[6])
ki_max = abs(g[7]) or 0.25
st["lean_i"] = float(np.clip(
st.get("lean_i", 0.0)
+ ki * st["odo"] * env.control_dt,
-ki_max, ki_max))
trim = (math.radians(lean_trim_deg(stance))
* (trim_s or 1.0) + st["lean_i"])
if stage in ("cruise", "turn", "slip", "bump",
"roam"):
# the stage's own speed loop: filtered speed and a
# ramped command, driving the lean target - the same
# law the episodes score
v_tau = abs(g[4]) or 0.10
_acc = max(0.3, abs(g[5]) or 1.2)
ph = (st["k"] * env.control_dt) % 12.0
if stage == "roam":
# FREE ROAM: the cockpit's own speed command,
# driving on the benches' baked balancing
want_v = float(np.clip(
b.get("roam_v", 0.0), -1.5, 1.5))
elif stage == "cruise":
want_v = (0.0 if ph < 1.5 else 0.8 if ph < 5.5
else 0.4 if ph < 7.5
else -0.5 if ph < 10.5 else 0.0)
elif stage == "turn":
want_v = 0.0 if ph < 1.5 else 0.5
else:
want_v = 0.0 if ph < 1.5 else 0.8
if stage == "slip":
if st.get("mu0") is None:
st["mu0"] = env.model.geom_friction[:, 0].copy()
on_ice = 5.0 <= ph < 7.5
env.model.geom_friction[:, 0] = st["mu0"] * (
0.25 if on_ice else 1.0)
st["on_ice"] = on_ice
st["v_cmd"] = st.get("v_cmd", 0.0) + float(np.clip(
want_v - st.get("v_cmd", 0.0),
-_acc * env.control_dt, _acc * env.control_dt))
st["v_f"] = st.get("v_f", 0.0) + (
env.control_dt / max(env.control_dt, v_tau)) * (
speed - st.get("v_f", 0.0))
st["oerr"] = (st.get("oerr", 0.0)
+ (st["v_f"] - st["v_cmd"])
* env.control_dt)
ref = CTRL.lean_reference(stn_g,
st["v_f"] - st["v_cmd"],
st["oerr"], clamp_g)
if st.get("reflex"):
ref += st["reflex"].update(
speed, st["v_cmd"], env.control_dt)
elif stage == "stop":
# Run up to speed, then stop as hard as it can, on a
# loop - the skill is only visible in use. The cruise
# RAMPS at 1.5 m/s^2 exactly as the scored episode
# does; the old step reference was a different law.
brake_g = g[0] if len(g) > 0 else 0.5
brake_c = g[1] if len(g) > 1 else 0.45
st["phase"] = st.get("phase", 0) + 1
if (st["phase"] % 320) < 170: # cruise
st["v_cmd"] = min(1.5, st.get("v_cmd", 0.0)
+ 1.5 * env.control_dt)
ref = max(-abs(clamp_g), min(abs(clamp_g),
stn_g * 0.5 * (speed - st["v_cmd"])))
else: # stop
st["v_cmd"] = 0.0
ref = max(-abs(brake_c), min(abs(brake_c),
brake_g * speed))
elif stage == "spin":
# the exam's law verbatim: NO odometry in the
# wheels (stillness doctrine) - the generic
# branch's homing term pegged the lean toward
# accumulated odo garbage after every spin
# (operator: 'a lean it cannot recover from')
ref = CTRL.lean_reference(stn_g, speed,
0.0, clamp_g)
if st.get("reflex"):
ref += st["reflex"].update(
speed, 0.0, env.control_dt)
elif stage == "jump":
# station hold, no odo (v56: no rolling start)
ref = CTRL.lean_reference(stn_g, speed,
0.0, clamp_g)
elif stage == "rotate" and len(g) >= 6:
# NO odometry in rotate's wheels (stillness
# doctrine): position belongs to the feet
ref = CTRL.lean_reference(stn_g, speed, 0.0,
clamp_g)
if st.get("reflex"):
ref += st["reflex"].update(
speed, 0.0, env.control_dt)
else:
# the benches carry no position memory (odo would
# lean them home, which the operator ruled out)
_odo = (0.0 if stage in ("balance", "balance_mid",
"balance_crouch")
else st["odo"])
ref = CTRL.lean_reference(stn_g, speed, _odo,
clamp_g)
if stage in ("spin", "payload") and st.get("reflex"):
ref += st["reflex"].update(
speed, 0.0, env.control_dt)
if stage == "recover":
# the brace the stage LEARNS in slots 8/9 - the demo
# used to unpack only g[:8] and silently dropped it
if st.get("brace") is None:
st["brace"] = CTRL.ShoveBrace(brace_g, brace_f)
st["brace"].gain, st["brace"].fade = brace_g, brace_f
ref += st["brace"].update(speed, 0.0, env.control_dt)
u = CTRL.wheel_command(kp_g, kd_g, lean, rate, ref, trim)
if stage == "jump":
u += st.get("j_catch_u", 0.0)
if stage in ("rise", "sit"):
st["uw_pair"] = None
if st.get("rise_hold"):
# lay/tuck phases: wheels position-locked
# PER WHEEL - the averaged shared command
# could not pin them and the robot slid
# backwards through the whole sequence
_rl = st.get("rlock") or [0.0, 0.0]
_uw = []
for wi, n2 in enumerate(
("l_wheel_j", "r_wheel_j")):
adr = env.model.jnt_qposadr[
mujoco_jid(env, n2)]
dof = env.model.jnt_dofadr[
mujoco_jid(env, n2)]
_uw.append(float(np.clip(
-3.0 * (float(env.data.qpos[adr])
- _rl[wi])
- 0.3 * float(env.data.qvel[dof]),
-1.0, 1.0)))
st["uw_pair"] = _uw
u = 0.0
elif st.get("rise_drive") is not None:
u = st["rise_drive"] # the rock-over push
if stage == "rotate":
# the chassis lives on a circle: wrap the lean
# error so 350 and -10 are neighbours, not a
# full-turn correction
_e = lean - ref - trim
_e = math.atan2(math.sin(_e), math.cos(_e))
u = float(np.clip(-kp_g * _e + kd_g * rate,
-1.0, 1.0))
if stage == "slip":
# the searched grip ceiling: on ice you ask for less
_cap = min(1.0, max(0.2, abs(g[6]) or 1.0))
u = float(np.clip(u, -_cap, _cap))
turn_t = 0.0
if stage == "spin" and len(g) >= 10:
gz_s = float(env.sensors.noisy(
"gyro", env.data.sensordata, env._raw)[2])
if b.get("spin_v") is not None:
# THE THROTTLE (operator's spec): a live
# proportional rate stick - slide for speed,
# centre for hold; the learned safe rate is
# the ceiling
_sv = float(np.clip(b["spin_v"], -1.0, 1.0))
_srm = min(4.0, max(0.2, abs(g[4])))
_wt = _sv * _srm
# ramp the rate in/out (3 rad/s^2): a
# stepped command kicks a translation the
# rotation smears into an ellipse
_wp = st.get("wz_sl", 0.0)
_wp += float(np.clip(_wt - _wp,
-3.0 * env.control_dt,
3.0 * env.control_dt))
st["wz_sl"] = _wp
wz_c = _wp
if st.get("sp_hold"):
wz_c = 0.0 # gate moving: balance only
turn_t = float(np.clip(
-abs(g[1]) * (wz_c - gz_s),
-min(0.5, max(0.02, abs(g[2]))),
min(0.5, max(0.02, abs(g[2])))))
turn_t = float(np.clip(
turn_t, -(1.0 - abs(u)), 1.0 - abs(u)))
st["h_est"] = 0.0
st["h_target"] = 0.0
st["sp_q"] = 0
_spin_done = True
else:
_spin_done = False
# indexed quarters: spin at the learned safe rate,
# decelerating as the remaining angle shrinks; next
# quarter once landed and quiet
st["h_est"] = st.get("h_est", 0.0) + gz_s * env.control_dt
st.setdefault("h_target", 0.0)
rem = st["h_target"] - st["h_est"]
if abs(rem) < math.radians(3.0) and abs(gz_s) < 0.25:
st["settled"] = st.get("settled", 0.0) + env.control_dt
else:
st["settled"] = 0.0
if _spin_done:
pass
elif st.get("settled", 0.0) >= 0.6:
# the operator's progression: 90, then 180,
# then the full 360, cycling
_dq = (90.0, 180.0, 360.0)[
st.get("sp_q", 0) % 3]
st["sp_q"] = st.get("sp_q", 0) + 1
st["h_target"] += math.radians(_dq)
st["settled"] = 0.0
if not _spin_done:
_sr = min(4.0, max(0.2, abs(g[4])))
if abs(rem) < math.radians(6.0) \
and abs(gz_s) < 0.4:
_sr = 0.3 # landed: balance, gentle trim
if st.get("sp_hold"):
_sr = 0.0 # gate moving: balance only
wz_c = float(np.clip(abs(g[3]) * rem,
-_sr, _sr))
turn_t = float(np.clip(
-abs(g[1]) * (wz_c - gz_s),
-min(0.5, max(0.02, abs(g[2]))),
min(0.5, max(0.02, abs(g[2])))))
turn_t = float(np.clip(
turn_t, -(1.0 - abs(u)), 1.0 - abs(u)))
elif stage == "roam":
# assist steering (honest label: the turn stage is
# not baked yet - when it is, its yaw law slots in
# here): small differential, balance keeps priority
_st_w = float(np.clip(b.get("roam_w", 0.0), -1, 1))
turn_t = float(np.clip(0.28 * _st_w,
-(1.0 - abs(u)),
1.0 - abs(u)))
elif stage == "turn":
yaw_rate = float(env.sensors.noisy(
"gyro", env.data.sensordata, env._raw)[2])
yaw_cmd = st.get("yce", 0.0)
k_yaw = abs(g[6]) or 0.35
t_max = min(0.5, max(0.02, abs(g[7]) or 0.30))
turn_t = float(np.clip(-k_yaw * (yaw_cmd - yaw_rate),
-t_max, t_max))
# balance first, turning with what is left
turn_t = float(np.clip(turn_t, -(1.0 - abs(u)),
1.0 - abs(u)))
st.setdefault("u_hist", []).append(abs(u))
del st["u_hist"][:-300]
LIVE["ctl"] = {
"u": round(u, 3), "lean_deg": round(math.degrees(lean), 2),
"ref_deg": round(math.degrees(ref + trim), 2),
"rate": round(rate, 3),
"sat_pct": round(100.0 * float(np.mean(
[1.0 if v > 0.95 else 0.0 for v in st["u_hist"]])), 1),
"mean_u": round(float(np.mean(st["u_hist"])), 3),
"brace": round(float(getattr(st.get("brace"), "brace", 0.0)), 4),
"v_cmd": round(float(st.get("v_cmd", 0.0)), 3),
"kg": st.get("loaded_kg"),
"on_ice": bool(st.get("on_ice")),
}
act = np.zeros(env.act_dim)
for j, ai in enumerate(env.act_idx):
if ai in legs and env.act_dim > 2:
act[j] = np.clip((st["leg"][legs.index(ai)]
- env.act_center[ai])
/ max(1e-9, env.act_span[ai]), -1, 1)
else:
nm3 = env.spec["actuators"][ai]["name"]
if st.get("uw_pair") is not None:
act[j] = st["uw_pair"][
0 if nm3.startswith("l") else 1]
else:
side = (turn_t if nm3.startswith("l")
else -turn_t)
act[j] = float(np.clip(u + side, -1, 1))
o, _, fell, trunc, _ = env.step(act)
LIVE["obs_now"] = o
# A shove on request - the broom. Applied once per call so
# holding the button does not integrate into a launch.
if st.get("pending_shove"):
env.data.qvel[0] += float(st.pop("pending_shove"))
elif stage == "recover" and b.get("auto_shove", True) \
and st["k"] % 100 == 0:
env.data.qvel[0] += 2.5 * (1 if (st["k"] // 100) % 2 else -1)
# Only restart if it actually FELL. Resetting on trunc as
# well made the demo jump back to the start every few
# seconds when the task's episode limit expired, which
# looks exactly like the robot running away or toppling.
# Comprehensive restart guard. Three ways a demo robot
# was lost for good: (a) a physics blow-up goes NaN, and
# every comparison with NaN is FALSE - neither the fall
# check nor any distance check can ever fire on a flung
# robot; (b) the 2m rule was level-only, so other stages
# had no boundary; (c) fell alone missed both.
_r = math.hypot(float(env.base_pos[0]),
float(env.base_pos[1]))
_lim = (60.0 if stage == "roam"
else 10.0 if stage in ("stop", "cruise", "turn",
"slip", "bump")
else 2.0 if stage == "level" else 3.0)
_bad = (not np.isfinite(env.data.qpos).all()
or not math.isfinite(_r) or _r > _lim)
if stage == "sit":
fell = bool(env.up_z < 0.25) # sideways only
if stage == "jump":
# flight and touchdown swing up_z by design -
# only a true tumble is a fall. A clean landing
# carries straight on balancing, so the energy
# trace survives and the next press is instant.
fell = bool(env.up_z < 0.35)
if stage == "rise":
# the lay and tuck put the FACE ON THE FLOOR by
# design - up_z is meaningless until the catch;
# only a topple after balance-on is a fall
fell = (st.get("rph", 0) >= 3 and env.up_z < 0.5)
if stage == "rotate":
# up_z is meaningless on the far side of the
# circle - a chassis hanging at 180 is a POSE, not
# a fall. Fallen = the legs collapsed, judged by
# hip-axis height against the commanded stance.
_hj = [mujoco_jid(env, n) for n in ("l_hip", "r_hip")]
_hz = 0.5 * (float(env.data.xanchor[_hj[0]][2])
+ float(env.data.xanchor[_hj[1]][2]))
fell = _hz < 0.55 * (stance / 1000.0)
if fell or _bad:
if stage == "jump":
print(f"JUMP-RESET fell={fell} bad={_bad} "
f"up_z={float(env.up_z):.3f} "
f"h={float(env.base_pos[2]):.3f} "
f"minh={env.min_h} minup={env.min_up}",
flush=True)
if stage == "sit":
print(f"SIT-RESET fell={fell} bad={_bad} "
f"up_z={float(env.up_z):.3f} r={_r:.2f} "
f"h={float(env.base_pos[2]):.3f} "
f"minh={env.min_h} minup={env.min_up}",
flush=True)
LIVE["obs_now"] = env.reset()
# a crash respawns STANDING - the rise demo must
# re-seat or a failed attempt restarts upright,
# and it seats RIGHT HERE so no standing frame
# ever flashes between attempts
st.update(odo=0.0, leg=None, k=0,
seated=False, rise_t=0.0)
if stage == "jump":
# a botched landing respawns standing with
# the sequence disarmed
st["jph"] = None
st["j_dir"] = None
st["j_st"] = 455.0
st["uw_pair"] = None
if stage == "sit":
# a crash mid-shutdown must give the power
# back and restart the descent from the top
if st.get("pwr_gain") is not None:
env.model.actuator_gainprm[:] = (
st["pwr_gain"])
env.model.actuator_biasprm[:] = (
st["pwr_bias"])
st["pwr_gain"] = None
st["pwr_bias"] = None
st["sph"] = None
if stage == "rise":
if str(TRAINER) not in sys.path:
sys.path.insert(0, str(TRAINER))
import learn_balance as _LBc
env.min_h = 0.05
_LBc.place_sitting(env)
_LBc.limp_settle(env, 0.6)
st["seated"] = True
# on the table, a restart PLACES the robot back on
# the flat plate in pose - reset alone leaves it at
# floor height inside the table geometry
if LIVE.get("world") == "table" and stage == "level":
_st_mm = float(np.clip(
400.0 + 100.0 * (g[8] if len(g) > 8 else 0.2),
340.0, 462.0))
_place_on_table(env, _st_mm)
break
if trunc:
env.step_i = 0 # a demo has no episode limit
act = None
elif b.get("mode") == "assist":
# Drive it like a real machine: the balance loop is always
# running off the sensors, and the sliders ask for leg angles
# and a road SPEED on top of it. The wheels are never commanded
# directly - asking for speed and letting the balancer decide
# the torque is the only thing that works, because a wheeled
# pendulum has to lean the wrong way before it can change speed.
deg = np.array(b.get("degrees", [0.0] * 4), dtype=float)
names = [a["name"] for a in env.spec["actuators"]]
legs = [i for i, a in enumerate(env.spec["actuators"])
if a.get("mode") == "position"]
whl = [i for i, a in enumerate(env.spec["actuators"])
if a.get("mode") == "torque"]
# The policy's action window is a NARROW band around a
# nominal crouch - hip +25, knee -50, span 40 - so a policy can
# only make small corrections. Driving by hand needs the whole
# joint, and with the narrow window a straight knee needed
# action 1.25 and clipped at 1.0: the robot could never stand
# up straight, sat permanently bent, and drifted from the
# offset centre of mass. Widen it for the hand-driven sim only.
for i in legs:
env.act_center[i] = 0.5 * (env.ctrl_lo[i] + env.ctrl_hi[i])
env.act_span[i] = 0.5 * (env.ctrl_hi[i] - env.ctrl_lo[i])
st = LIVE.setdefault("assist", {})
if b.get("zero") or "leg" not in st:
st["leg"] = np.array([float(env.data.qpos[
env.model.jnt_qposadr[mujoco_jid(env, names[i][:-2])]])
for i in legs])
st["x_ref"] = float(env.base_pos[0])
st["v_cmd"] = 0.0
st["odo"] = 0.0
st["branch"] = -1.0
st["swing"] = False
# These are the numbers you confirmed by hand in the bench:
# kp 12, kd 2, station keeping 1.95. Nothing clever on top -
# every "improvement" I tried (a wider lean clamp, a world-frame
# centre-of-mass feedforward, a trim integrator) made it fall
# over, because this loop is already tuned and its lean clamp is
# load-bearing, not incidental.
kp = float(b.get("kp", 12.0)); kd = float(b.get("kd", 2.0))
# 1.95 was right before the lean trim existed: the station loop had to
# fight the whole crouch offset on its own. Now the trim supplies
# that directly, and a gain this high just oscillates - the wheels
# buzz back and forth at +-85rpm forever after the robot has
# stopped, which is what "the wheels keep going" was. 0.40 cuts
# that to 3.5rpm and holds position at least as well.
station = float(b.get("station", 1.0))
# 6.9 deg, not 4.6. A balancing robot's acceleration IS its lean
# (a = g*tan(lean)), so the cap - not motor torque - is what
# makes it feel sluggish. Measured: 4.6 deg reaches 0.8 m/s in
# 1.62s and survives a 2.0 m/s shove; 6.9 does it in 1.51s and
# still survives 2.0. Past about 8 deg it gets quicker but
# starts dropping shoves, so this is the free part of the trade.
lean_max = float(b.get("lean_max", 0.12))
v_max = float(b.get("v_max", 1.0))
leg_rate = math.radians(float(b.get("leg_rate_deg", 60.0)))
swing_rate = math.radians(float(b.get("swing_rate_deg", 400.0)))
acc = float(b.get("accel", 1.2))
# Stance height, by your maths. Two 200mm segments: fold the
# knee by q and the leg spans 2*200*cos(q/2) from hip to axle.
# Driving the hip to +q/2 and the knee to -q puts the ankle
# exactly under the hip - sin(q/2) + sin(q/2 - q) = 0 - so the
# wheel stays beneath the hip at every height.
turn0 = float(b.get("turn", 0.0))
k_roll = float(b.get("k_roll", 0.0))
k_rolld = float(b.get("k_rolld", 0.0))
roll_ff = float(b.get("roll_ff", 0.0))
want_v = np.clip(float(b.get("drive", 0.0)) / 100.0, -1, 1) * v_max
shift_k = float(b.get("shift_k", 0.30))
shift_max = float(b.get("shift_max", 0.10))
have_stance = b.get("stance_mm") is not None
h = float(b.get("stance_mm", 460.0)) / 1000.0
if not have_stance:
want_leg = np.radians(deg[:len(legs)])
want_leg = np.clip(want_leg,
[env.ctrl_lo[i] for i in legs],
[env.ctrl_hi[i] for i in legs])
lean_bias = np.radians(float(b.get("lean_deg", 0.0)))
# Turn is a RATE command, closed on the gyro. It used to be a
# plain torque difference, which is a constant yaw ACCELERATION
# with nothing to damp it. Sign is measured, not assumed: a
# differential of +0.30 on (left = u+t, right = u-t) yaws the
# robot NEGATIVE, so the error term is negated.
yaw_cmd = np.clip(float(b.get("turn", 0.0)) / 100.0, -1, 1) \
* float(b.get("yaw_max", 1.2))
k_yaw = float(b.get("k_yaw", 0.35))
# Low-pass the measured speed before it drives the lean
# target. Raw wheel speed contains the limit cycle itself, so
# feeding it back unfiltered closes a positive loop: the lean
# target saturates at its clamp, flips sign every step, and
# the wheels chatter at full torque while the robot stands
# perfectly still - measured 2.22 Nm mean where holding still
# needs about 0.05. assist_lab documented the fix (v_tau) and
# this copy of the loop never had it.
v_tau = float(b.get("v_tau", 0.10))
sat = 0
for _ in range(int(b.get("substeps", 3))):
dt = env.control_dt
# Everything that feeds the wheels runs at the full
# 100 Hz. The foot shift, roll legs, trim and yaw loop
# used to be computed once per HTTP call (~33 Hz with
# substeps=3) - ~30 ms of extra phase lag inside a closed
# pitch loop, present here and absent from auto-balance.
up = env.est_up # sensors, not truth
pitch = math.atan2(-up[0], max(1e-6, up[2]))
g3 = env.sensors.noisy("gyro", env.data.sensordata,
env._raw)
rate = float(g3[1])
roll_rate = float(g3[0])
yaw_rate = float(g3[2])
v_raw = 0.5 * sum(float(env.data.qvel[env.model.jnt_dofadr[
mujoco_jid(env, n)]])
for n in ("l_wheel_j", "r_wheel_j")) * 0.0625
if "v_f" not in st:
st["v_f"] = v_raw
st["v_f"] += (dt / max(dt, v_tau)) * (v_raw - st["v_f"])
v = st["v_f"]
if have_stance:
# Roll control, using the legs. The wheels share one
# axle and can do nothing about roll; one leg longer
# than the other leans the body like a motorbike.
roll = math.atan2(up[1], max(1e-6, up[2]))
dh = (k_roll * roll + k_rolld * roll_rate
+ roll_ff * (turn0 / 100.0))
dh = max(-0.06, min(0.06, dh))
# Feet TOWARD the fall (sign lives in
# controller.foot_shift - this file's copy negated it,
# turning a measured 24.9->6.7 deg recovery reflex
# into a delayed positive feedback path from pitch
# into the contact patch, on by default while driving)
dx = CTRL.foot_shift(shift_k, pitch, shift_max)
# Posture follows DIRECTION OF TRAVEL; swap only from
# a standstill, legs thrown across at swing speed
# with the drive held at zero until they arrive.
want_br = st.get("branch", -1.0)
if want_v > 0.05:
want_br = -1.0
elif want_v < -0.05:
want_br = +1.0
st["pending_br"] = want_br
if want_br != st.get("branch", -1.0) and abs(v) < 0.10:
st["swing"] = True
st["branch"] = want_br
br = st.get("branch", -1.0)
pairs = [leg_ik(max(0.20, min(0.462, h + side * dh)),
dx, br)
for side in (+1.0, -1.0)]
want_leg = np.array([pairs[0][0], pairs[0][1],
pairs[1][0],
pairs[1][1]][:len(legs)])
want_leg = np.clip(want_leg,
[env.ctrl_lo[i] for i in legs],
[env.ctrl_hi[i] for i in legs])
br = st.get("branch", -1.0)
# The trim FLIPS with the posture: knee-back needs
# -4.95 deg at 400mm, knee-forward +5.18.
trim = (math.radians(lean_trim_deg(
float(b.get("stance_mm", 460.0))))
* (1.0 if br < 0 else -1.0)
) if b.get("trim", True) else 0.0
# Move the joints TOGETHER (one scale factor for the whole
# move keeps the axle under the hip the entire way).
delta = want_leg - st["leg"]
far = float(np.max(np.abs(delta))) if delta.size else 0.0
swinging = bool(st.get("swing")) and far > math.radians(4.0)
st["swing"] = swinging
rate_now = swing_rate if swinging else leg_rate
if far > 1e-9:
st["leg"] = st["leg"] + delta * (
min(far, rate_now * dt) / far)
pending = st.get("pending_br", br) != st.get("branch", br)
v_target = 0.0 if (swinging or pending) else want_v
st["v_cmd"] += float(np.clip(v_target - st["v_cmd"],
-acc * dt, acc * dt))
# odometry of the SPEED ERROR, so asking for speed does not
# wind the station-keeping term up against you
st["odo"] = st.get("odo", 0.0) + (v - st["v_cmd"]) * dt
lean = max(-lean_max, min(lean_max, station * (
0.5 * (v - st["v_cmd"]) + 0.1 * st["odo"])))
u = CTRL.wheel_command(kp, kd, pitch, rate,
lean, lean_bias + trim)
if abs(u) >= 0.999:
sat += 1
turn = float(np.clip(-k_yaw * (yaw_cmd - yaw_rate),
-0.30, 0.30))
# Balance first, turning with whatever is left: without
# this reserve (assist_lab has it; this copy did not) a
# turn differential clips one wheel exactly when the
# balancer needs a big surge, and the robot cannot drive
# hard enough to catch itself.
turn = float(np.clip(turn, -(1.0 - abs(u)),
1.0 - abs(u)))
LIVE["dbg"] = {"pitch": round(math.degrees(pitch), 2),
"lean": round(math.degrees(lean), 2),
"trim": round(math.degrees(trim), 2),
"rate": round(rate, 3),
"v": round(v, 3), "v_cmd": round(st["v_cmd"], 3),
"odo": round(st["odo"], 3),
"kp_t": round(-kp * (pitch - lean - lean_bias - trim), 3),
"kd_t": round(kd * rate, 3),
"u": round(u, 3), "turn": round(turn, 3)}
act = np.zeros(env.act_dim)
for j, ai in enumerate(env.act_idx):
if ai in whl:
side = turn if names[ai].startswith("l") else -turn
act[j] = np.clip(u + side, -1, 1)
else:
k = legs.index(ai)
act[j] = np.clip((st["leg"][k] - env.act_center[ai])
/ max(1e-9, env.act_span[ai]), -1, 1)
if LIVE["rec"] and LIVE.get("obs_now") is not None:
LIVE["obs"].append(list(map(float, LIVE["obs_now"])))
LIVE["act"].append(list(map(float, act)))
o, _, fell, trunc, _ = env.step(act)
LIVE["obs_now"] = o
LIVE["assist_sat"] = sat
act = None
elif b.get("mode") == "angle":
# Calibration drives real joint angles over the FULL travel,
# not the narrow window the policy commands around the nominal
# stance. Otherwise the sliders cover only part of the joint.
deg = np.array(b.get("degrees", [0.0] * env.act_dim), dtype=float)
ctrl = np.zeros(env.act_dim)
for i in range(env.act_dim):
# A wheel slider is a percentage of available torque; a leg
# slider is an angle. Deciding by act_span was wrong - it is
# under 1.0 for both, so wheels were read as radians and
# saturated at about 19 percent.
if env.spec["actuators"][i].get("mode") == "torque":
sgn = -1.0 if env.spec["actuators"][i].get("reversed") else 1.0
ctrl[i] = np.clip(deg[i] / 100.0, -1, 1) * env.ctrl_hi[i] * sgn
else:
ctrl[i] = np.clip(math.radians(deg[i]),
env.ctrl_lo[i], env.ctrl_hi[i])
# A free-spinning wheel cannot hold the robot still: nothing in
# the model damps horizontal speed, so any nudge from the spawn
# drop rolls away forever. The real motors would not allow that
# - a FOC driver holds the rotor against a load. With the drive
# slider at zero, brake to a stop within the motor's torque.
brake = bool(b.get("brake", True))
wheels = [i for i, a in enumerate(env.spec["actuators"])
if a.get("mode") == "torque"]
idle = [i for i in wheels if abs(deg[i]) < 1e-6]
for _ in range(int(b.get("substeps", 3))):
if brake and idle:
for i in idle:
w = float(env.data.qvel[env.act_dof[i]])
ctrl[i] = float(np.clip(-0.05 * w, -env.ctrl_hi[i],
env.ctrl_hi[i]))
env.data.ctrl[:] = ctrl
for _ in range(env.n_sub):
mujoco_step(env)
env.t += env.control_dt
env._cache_truth(); env._est_update()
act = None
else:
act = np.array(b.get("action", [0.0] * env.act_dim), dtype=float)
act = np.clip(act, -1.0, 1.0)[:env.act_dim]
# "Tethered" holds the body LEVEL only. Height stays free so the
# legs can lift and lower it, and fore/aft stays free so it can
# actually drive - otherwise you can never see whether the wheels
# move the robot.
if b.get("auto_balance"):
kp = float(b.get("kp", 12.0)); kd = float(b.get("kd", 2.0))
# 1.0. Measured on the reverted 1:1 wheels: 0.40 drifts 846mm in 9s,
# 1.0 halves that for 0.23 Nm, and 1.95 buys nothing more while
# costing torque and pitch precision.
st = float(b.get("station", 1.0))
r_wheel = 0.0625
for _ in range(int(b.get("substeps", 3))):
up = env.est_up
pitch = math.atan2(-up[0], max(1e-6, up[2]))
rate = float(env.sensors.noisy("gyro", env.data.sensordata, env._raw)[1])
wv = 0.5 * (env.data.qvel[env.model.jnt_dofadr[mujoco_jid(env, "l_wheel_j")]]
+ env.data.qvel[env.model.jnt_dofadr[mujoco_jid(env, "r_wheel_j")]])
speed = wv * r_wheel
LIVE["odo"] = LIVE.get("odo", 0.0) + speed * env.control_dt
# cascade: feeding speed straight into torque destabilises
# (to stop, the robot must first lean the other way). Speed
# and odometry pick a small LEAN TARGET; balance chases it.
lean_ref = CTRL.lean_reference(st, speed, LIVE["odo"], 0.08)
u = CTRL.wheel_command(kp, kd, pitch, rate, lean_ref)
# both wheels the same way: they share a physical axis
env.step(np.array([0.0, 0.0, 0.0, 0.0, u, u])[env.act_idx]
if env.act_dim == 6 else np.array([u, u]))
act = None
if act is not None:
for _ in range(int(b.get("substeps", 3))):
# record what the robot could see and what you did, so the
# pair can be cloned later
if LIVE["rec"] and LIVE.get("obs_now") is not None:
LIVE["obs"].append(list(map(float, LIVE["obs_now"])))
LIVE["act"].append(list(map(float, act)))
o, _, fell, trunc, _ = env.step(act)
LIVE["obs_now"] = o
if fell or trunc:
LIVE["obs_now"] = env.reset()
break
# Black box. Every live step goes into a rolling buffer so that when
# something happens once, in front of you, we can go and look at it
# instead of trying to reproduce it. Cheap enough to always be on.
try:
_qa = float(np.max(np.abs(env.data.qacc)))
_cf = 0.0; _self = None
for _i in range(env.data.ncon):
_f6 = np.zeros(6)
mujoco.mj_contactForce(env.model, env.data, _i, _f6)
if _f6[0] > _cf:
_cf = float(_f6[0])
_c = env.data.contact[_i]
_b1 = mujoco.mj_id2name(env.model, mujoco.mjtObj.mjOBJ_BODY,
env.model.geom_bodyid[_c.geom1])
_b2 = mujoco.mj_id2name(env.model, mujoco.mjtObj.mjOBJ_BODY,
env.model.geom_bodyid[_c.geom2])
_self = f"{_b1}|{_b2}"
BLACKBOX.append({
"wall": round(time.time(), 3), "t": round(env.t, 3),
"cmd": [round(float(v), 2) for v in b.get("degrees", [])],
# assist commands, and ROLL. The spin that fell on its face
# could not be diagnosed from this buffer because only
# pitch-ish up_z was here: a robot thrown outward in roll by
# a hard turn looks identical to one falling forward.
"drive": b.get("drive"), "turn": b.get("turn"),
"stance": b.get("stance_mm"), "lean_cmd": b.get("lean_deg"),
"pitch": round(math.degrees(math.atan2(
-env.est_up[0], max(1e-6, env.est_up[2]))), 2),
"roll": round(math.degrees(math.atan2(
env.est_up[1], max(1e-6, env.est_up[2]))), 2),
"mode": b.get("mode"), "auto": bool(b.get("auto_balance")),
"tether": b.get("tether"), "brake": bool(b.get("brake", True)),
"tau": [round(float(v), 3) for v in env.data.actuator_force],
"q": [round(float(env.data.qpos[i]), 4) for i in env.leg_joint_ids],
"qd": [round(float(env.data.qvel[env.model.jnt_dofadr[
mujoco_jid(env, n)]]), 2) for n in ("l_wheel_j", "r_wheel_j")],
"h": round(env.height, 4), "up": round(env.up_z, 3),
"x": round(float(env.base_pos[0]), 4),
"qacc": round(_qa, 1), "cf": round(_cf, 1), "con": _self,
"ncon": int(env.data.ncon),
})
except Exception:
pass
return self._send({
"frame": env.frame(),
"t": round(env.t, 3),
"height": round(env.height, 4),
"up_z": round(env.up_z, 4),
"pitch_deg": round(math.degrees(math.atan2(-env.est_up[0], max(1e-6, env.est_up[2]))), 2),
"roll_deg": round(math.degrees(math.atan2(env.est_up[1], max(1e-6, env.est_up[2]))), 2),
"tau": [round(float(v), 4) for v in env.data.actuator_force],
"qpos": [round(float(env.data.qpos[i]), 4) for i in env.leg_joint_ids],
"fell": bool(env.up_z < env.min_up or env.height < env.min_h),
"recording": LIVE["rec"],
"demo_steps": len(LIVE["obs"]),
"task": LIVE["task"],
"act_dim": env.act_dim,
"x": round(float(env.base_pos[0]), 4),
"y": round(float(env.base_pos[1]), 4),
"vx": round(float(env.vel_local[0]), 3),
"dbg": LIVE.get("dbg"),
"ctl": LIVE.get("ctl"),
"slope_deg": round(LIVE.get("gainstate", {}).get("slope_deg", 0.0), 2),
"spill_deg": round(LIVE.get("gainstate", {}).get("spill_deg", 0.0), 2),
"wheel_rpm": [round(float(env.data.qvel[
env.model.jnt_dofadr[mujoco_jid(env, n)]]) * 60 / (2*math.pi), 1)
for n in ("l_wheel_j", "r_wheel_j")],
# what the body actually senses, at this instant - the scope
# panel draws these, because every fault so far was invisible
# until someone happened to measure the right signal
"gyro": [round(float(v), 4) for v in env.sensors.noisy(
"gyro", env.data.sensordata, env._raw)],
"accel": [round(float(v), 3) for v in env.sensors.noisy(
"accel", env.data.sensordata, env._raw)],
})
def _clone(self, b):
"""Turn recorded demonstrations into a starting policy."""
name = b.get("name", "")
if not SAFE_NAME.match(name):
raise ValueError("bad demonstration name")
if _busy():
return self._send({"error": "stop the running job first"}, 409)
run = b.get("run") or "johnny6"
task = b.get("task") or "balance"
out = ROOT / "runs" / f"{run}__{task}"
argv = ["demo.py", "--clone", name, "--task", task,
"--out", str(out), "--epochs", str(int(b.get("epochs", 300)))]
j = _launch("clone", argv, label=f"clone {name} -> {task}")
self._send(_job_view(j))
def _reset_step(self, b):
"""Throw away a stage's progress so it trains from scratch."""
run = b.get("run") or "johnny6"
idx = int(b.get("index", 1))
if not SAFE_NAME.match(run):
raise ValueError("bad run name")
if _busy():
return self._send({"error": "stop the running job first"}, 409)
cp = ROOT / "runs" / run / "course.json"
if cp.exists():
state = json.load(open(cp))
st = state["stages"][idx - 1]
st.update({"status": "pending", "exam": None, "policy": None})
json.dump(state, open(cp, "w"), indent=1)
import shutil
shutil.rmtree(ROOT / "runs" / st["run"], ignore_errors=True)
self._send({"ok": True})
def _stop(self, b):
j = JOBS.get(b.get("id"))
if not j:
raise ValueError("no such job")
if _job_alive(j):
try:
os.killpg(os.getpgid(j["pid"]), signal.SIGTERM)
except OSError:
pass
self._send(_job_view(j))
def log_message(self, fmt, *a):
# args are not always strings (send_error passes an HTTPStatus,
# which crashed the handler thread on every 404)
if "/api/" in (str(a[0]) if a else ""):
return
super().log_message(fmt, *a)
def main():
port = int(sys.argv[1]) if len(sys.argv) > 1 else 8901
handler = partial(Handler, directory=str(ROOT))
srv = ThreadingHTTPServer(("127.0.0.1", port), handler)
_load_jobs()
adopted = sum(1 for j in JOBS.values() if _job_alive(j))
# Keep the robot present in Ghrups (online heartbeat) whenever the bench is up.
threading.Thread(target=_ghrups_heartbeat_loop, daemon=True).start()
print(f"robot bench: http://127.0.0.1:{port}/bench/index.html")
print(f"{CORES} cores; {adopted} job(s) still running from a previous server")
try:
srv.serve_forever()
except KeyboardInterrupt:
# Leave detached jobs running; they are re-adopted on next start.
_save_jobs()
if __name__ == "__main__":
main()