Download training/scripts/walk_loop.py from kenpeter123/twoleg: direct link, hf CLI and curl.
- Browser
- Download file 20.8 kB
-
https://huggingface.co/kenpeter123/twoleg/resolve/main/training/scripts/walk_loop.py
- Command line
-
hf download hf://kenpeter123/twoleg/training/scripts/walk_loop.py
-
curl -L -o walk_loop.py https://huggingface.co/kenpeter123/twoleg/resolve/main/training/scripts/walk_loop.py
20.8 kB
| """Reflective walk loop: train, watch the end of the video, reflect, small change, repeat. | |
| This is the agentic loop for /pp (TwoLeg biped). Each round: | |
| 1. resume-train a bounded iteration budget from the current checkpoint, | |
| 2. render a rollout and run scripts/verify_walk.py (which now judges the | |
| LAST frames: upright + both feet sharing load = walking), | |
| 3. if WALKS -> STOP (success), | |
| 4. if NO-WALK -> the REFLECTOR reads the verdict + end-frame pixels, writes | |
| ONE small reversible change as a new strategy, and the loop repeats. | |
| There is no fixed strategy table and no STRATEGIES_EXHAUSTED stop: the loop | |
| only ends on WALKS or a hard training failure. The reflector is the brain; it | |
| must keep each change small (one reward term, one servo sign, one hyper) so | |
| the search stays local and reversible. State persists in | |
| .audit/walk_loop_state.json so a later /pp call resumes. | |
| uv run python scripts/walk_loop.py --rounds 1 | |
| uv run python scripts/walk_loop.py --reflect-only --from-run <r> --from-checkpoint <pt> | |
| """ | |
| import argparse | |
| import json | |
| import os | |
| import subprocess | |
| import sys | |
| import time | |
| HERE = os.path.dirname(os.path.abspath(__file__)) | |
| TRAINING = os.path.dirname(HERE) | |
| REPO = os.path.dirname(TRAINING) | |
| AUDIT = os.path.join(REPO, ".audit") | |
| STATE = os.path.join(AUDIT, "walk_loop_state.json") | |
| STRAT = os.path.join(AUDIT, "walk_loop_strategies.jsonl") # append-only reflector log | |
| LOG_TSV = os.path.join(AUDIT, "twoleg-walk.tsv") | |
| LOG_SH = ("/home/kenpeter/work/pp/.opencode/skills/show-me-your-work/scripts/log.sh") | |
| TASK = "Mjlab-Velocity-Flat-TwoLeg" | |
| sys.path.insert(0, HERE) | |
| import loop_changes # whitelisted, git-committed, test-gated code edits | |
| # Reflection registry lives in code as the FIRST hypothesis; after that the | |
| # reflector appends. Seed with the standing open questions so round 1 has a | |
| # direction if state is fresh. | |
| SEED_STRATEGY = { | |
| "name": "fresh-symmetry-gate-1000", "kind": "resume-train", "iters": 1000, "envs": 4096, | |
| "why": "H12: last verdict left_duty 0.0 / right_duty 0.936 -> collapses onto " | |
| "right leg, left leg does nothing. Hypothesis: leg servos asymmetric " | |
| "(sign/scale mismatch or unmirrored position target). Add a left/right " | |
| "duty-balance reward term so both legs must bear load; predict symmetric " | |
| "stance and stepping. Revert if it flails.", | |
| } | |
| def sh(cmd, **kw): | |
| return subprocess.run(cmd, capture_output=True, text=True, **kw) | |
| def log_row(phase, decision, why, evidence, result): | |
| sh(["bash", LOG_SH, LOG_TSV, phase, decision, why, evidence, result]) | |
| def load_state(): | |
| if os.path.exists(STATE): | |
| with open(STATE) as fh: | |
| return json.load(fh) | |
| return {"rounds_done": 0, "run": None, "checkpoint": None, | |
| "last_verdict": None, "last_video": None} | |
| def save_state(s): | |
| os.makedirs(AUDIT, exist_ok=True) | |
| with open(STATE, "w") as fh: | |
| json.dump(s, fh, indent=1) | |
| def append_strategy(st): | |
| os.makedirs(AUDIT, exist_ok=True) | |
| with open(STRAT, "a") as fh: | |
| fh.write(json.dumps(st) + "\n") | |
| def run_dirs(): | |
| root = os.path.join(TRAINING, "logs", "rsl_rl", "velocity") | |
| return sorted(d for d in os.listdir(root) | |
| if os.path.isdir(os.path.join(root, d))) | |
| def newest_model(run): | |
| d = os.path.join(TRAINING, "logs", "rsl_rl", "velocity", run) | |
| if not os.path.isdir(d): | |
| return None | |
| pts = sorted(f for f in os.listdir(d) | |
| if f.startswith("model_") and f.endswith(".pt")) | |
| if not pts: | |
| return None | |
| return max(pts, key=lambda f: os.path.getmtime(os.path.join(d, f))) | |
| def train_round(strategy, run, ckpt, fresh=False): | |
| """Launch a bounded training run detached, poll until it exits. Returns the | |
| new run dir, or None on failure. | |
| fresh=True -> train from random init (NO --agent.resume / load flags). | |
| fresh=False -> resume from (run, ckpt). ANY run we launch is always | |
| recorded in state with its newest checkpoint, so it can be resumed later. | |
| """ | |
| iters, envs = strategy["iters"], strategy["envs"] | |
| before = set(run_dirs()) | |
| logf = f"/tmp/opencode/loop_{strategy['name']}.log" | |
| if os.path.exists(logf): | |
| os.remove(logf) | |
| resume_part = "" | |
| if not fresh: | |
| resume_part = (f"--agent.resume True --agent.load-run {run} " | |
| f"--agent.load-checkpoint {ckpt} ") | |
| # xvfb-run wraps train: --video True needs a GL context, and this host has | |
| # no real display. Without it the child dies at renderer init (caught as | |
| # "no new run dir appeared"). | |
| cmd = (f"setsid nohup xvfb-run -a -s \"-screen 0 800x600x24\" " | |
| f"uv run --no-sync train {TASK} " | |
| f"--env.scene.num-envs {envs} --agent.max-iterations {iters} " | |
| f"{resume_part}--video True " | |
| f"> {logf} 2>&1 < /dev/null &") | |
| subprocess.Popen(cmd, shell=True, cwd=TRAINING, | |
| stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) | |
| # Generous wall-clock budget: ~1.0s/iter on this 150W-capped 4070 Ti plus | |
| # startup (xvfb+uv+cuda) slack. Never pkill a healthy run that is still | |
| # making progress -- only give up if the proc actually died. The earlier | |
| # iters*2.0+600 budget was too tight and killed a resume at iter 3255. | |
| deadline = time.time() + iters * 1.0 + 1800.0 | |
| launched = False | |
| while time.time() < deadline: | |
| time.sleep(20) | |
| r = sh(["pgrep", "-f", "train Mjlab"]) | |
| alive = bool(r.stdout.strip()) | |
| if alive: | |
| launched = True | |
| elif launched: | |
| break # proc exited on its own -> run finished (or crashed) | |
| if time.time() >= deadline and launched: | |
| # Only reached if the proc is STILL alive past budget. Log but do not | |
| # kill; let it keep training and pick up the newer checkpoint. | |
| print(f"train round still running past budget ({iters} iters); " | |
| f"continuing to poll for checkpoint") | |
| after = set(run_dirs()) | |
| new = sorted(after - before) | |
| if not new: | |
| print("no new run dir appeared; check", logf) | |
| return None | |
| print(f"round done, new run dir: {new[-1]}") | |
| return new[-1] | |
| def verify(run, ckpt): | |
| """Run the metric verifier (headless, no GL) which reads the same contact | |
| duties + last-frame altitude the video verifier used. Returns (passed, sidecar). | |
| Returns (passed, sidecar). Returns a dict with 'verdict' in sidecar.""" | |
| if not ckpt: | |
| raise RuntimeError(f"verify: no checkpoint in run {run} (training died?)") | |
| ckpt_path = os.path.join(TRAINING, "logs", "rsl_rl", "velocity", run, ckpt) | |
| if not os.path.exists(ckpt_path): | |
| raise RuntimeError(f"verify: checkpoint missing: {ckpt_path}") | |
| env = dict(os.environ) # no MUJOCO_GL needed (verify_metrics has no render) | |
| sidecar_path = os.path.join(TRAINING, "logs", "rsl_rl", "velocity", run, | |
| "verify_metrics_cmd0.10.json") | |
| before = os.path.getmtime(sidecar_path) if os.path.exists(sidecar_path) else 0.0 | |
| r = subprocess.run( | |
| ["uv", "run", "--no-sync", "python", "scripts/verify_metrics.py", | |
| "--checkpoint", ckpt_path, "--command", "0.1", | |
| "--steps", "250", "--envs", "4", "--device", "cuda:0"], | |
| cwd=TRAINING, capture_output=True, text=True, env=env, timeout=300) | |
| print(r.stdout[-1500:] if r.stdout else "") | |
| if r.stderr: | |
| print("ERR:", r.stderr[-800:]) | |
| if not os.path.exists(sidecar_path): | |
| raise RuntimeError(f"verify: no sidecar written for {ckpt}") | |
| if os.path.getmtime(sidecar_path) <= before: | |
| raise RuntimeError( | |
| f"verify: sidecar is stale (not rewritten for {ckpt}); refusing to " | |
| "gate on an older checkpoint's numbers") | |
| with open(sidecar_path) as fh: | |
| sidecar = json.load(fh) | |
| REQUIRED = ("verdict", "valid_frac", "left_duty_mean", "right_duty_mean", | |
| "speed_mean", "last_frame_alt_mean", "switch_hz_mean", | |
| "double_support_frac", "knee_flex_min_rad", "hip_swing_min_rad", | |
| "stance_knee_bend_rad", "articulated_frac", "rolling_pass_frac") | |
| missing = [k for k in REQUIRED if k not in sidecar] | |
| if missing: | |
| raise RuntimeError(f"verify: sidecar missing keys {missing}; the gate " | |
| f"would read defaults, not measurements") | |
| return sidecar.get("verdict") == "WALKS", sidecar | |
| def reflect(state): | |
| """The brain. Read the last verdict + end-frame pixels, propose ONE small | |
| reversible change. Returns a strategy dict. Never invents a rewrite.""" | |
| v = state.get("last_verdict") or {} | |
| reasons = "; ".join(v.get("reasons", [])) or "no verdict" | |
| left = v.get("left_duty", "?") | |
| right = v.get("right_duty", "?") | |
| speed = v.get("speed", "?") | |
| last_min = v.get("last_min_duty", "?") | |
| last_alt = v.get("last_frame_alt", "?") | |
| video = state.get("last_video") or v.get("video") or "" | |
| round_n = state["rounds_done"] + 1 | |
| prompt = ( | |
| f"You are the reflection brain for a TwoLeg biped RL training loop. " | |
| f"Round {round_n} FAILED to walk.\n\n" | |
| f"VERIFIER VERDICT (last-frames aware):\n" | |
| f" reasons: {reasons}\n" | |
| f" duty left={left} right={right} speed={speed}\n" | |
| f" LAST-frames min duty={last_min} last frame_alt={last_alt}\n\n" | |
| f"WATCH THE VIDEO before deciding: {video}\n" | |
| f"Open it and look at the FINAL frames. Is the robot upright? Are BOTH " | |
| f"feet on the ground bearing load, or did it collapse onto one leg / fall? " | |
| f"Describe what the last frames actually show, then explain the most " | |
| f"likely cause.\n\n" | |
| f"Then propose EXACTLY ONE small, reversible change to make it walk on " | |
| f"two legs (no rewrites):\n" | |
| f" - add/adjust a reward term (duty-balance so both legs share load, " | |
| f"foot-contact alternation, torso-upright bonus)\n" | |
| f" - fix a leg servo sign/scale or mirror bug in the reference XML\n" | |
| f" - adjust one hyper (learning rate, iters, action clip)\n" | |
| f"Return JSON only: {{\"name\": str, \"kind\": \"resume-train\", \"iters\": 2000, " | |
| f"\"envs\": 2048, \"why\": str, \"change\": str}}" | |
| ) | |
| # Use the running opencode instance in poteto mode to reflect + watch the | |
| # end-frame video (it owns the model/XML context). Pin muse-spark-1.3 | |
| # (free) per direction. If opencode is unavailable, fall back to a | |
| # duty-balance nudge. NOTE: opencode can hang (keeps pipe open past | |
| # timeout) so we launch it in its own process group and killpg hard. | |
| import subprocess as _sp | |
| try: | |
| proc = _sp.Popen( | |
| ["opencode", "run", "--model", | |
| "opencode/muse-spark-1.3-contributor-free", prompt], | |
| cwd="/home/kenpeter/work/pp", | |
| stdout=_sp.PIPE, stderr=_sp.PIPE, text=True, | |
| start_new_session=True) | |
| try: | |
| out, err = proc.communicate(timeout=180) | |
| except _sp.TimeoutExpired: | |
| proc.kill(); _sp.run(["pkill", "-9", "-f", "opencode run"], | |
| capture_output=True) | |
| raise RuntimeError("opencode timed out after 180s") | |
| txt = (out or "") + (err or "") | |
| j = txt[txt.find("{"):txt.rfind("}")+1] | |
| st = json.loads(j) | |
| st.setdefault("iters", 1000); st.setdefault("envs", 4096) | |
| st.setdefault("kind", "resume-train") | |
| return st | |
| except Exception as e: | |
| print(f"reflect: opencode unavailable ({e}); fallback duty-balance") | |
| return {"name": f"duty-balance-{round_n}", "kind": "resume-train", | |
| "iters": 1000, "envs": 4096, | |
| "why": f"fallback: still collapsing (last min duty {last_min}). " | |
| f"Add duty-balance reward so both legs bear load.", | |
| "change": "add left/right duty-balance term to reward"} | |
| def decide(state, new_run, new_ckpt): | |
| """Agentic loop-condition. The GROUND TRUTH is the rigorous verify_metrics.py | |
| verdict (human-gait gate: upright + both feet load + signed speed + switch_hz | |
| + double-support + knee/hip articulation + smoothness). We MUST trust it, | |
| because a weaker local gate here already false-declared WALKS on a symmetric | |
| squat-freeze (upright + both feet down + tiny motion) and stopped the loop. | |
| - STOP if verdict == "WALKS" (verify_metrics said >=60% of envs pass the | |
| full human-gait gate). | |
| - CONTINUE otherwise -- the rigorous gate already encoded the evidence; we do | |
| NOT re-derive a weaker local gate that can be gamed. | |
| """ | |
| info = state.get("last_verdict") or {} | |
| verdict = info.get("verdict", "NO-WALK") | |
| if verdict == "WALKS": | |
| # Cross-check the evidence the verifier already measured (defense in | |
| # depth, but the verifier's verdict is authoritative). | |
| fa = info.get("last_frame_alt_mean", -1.0) | |
| l_d = info.get("left_duty_mean", 0.0) | |
| r_d = info.get("right_duty_mean", 0.0) | |
| sp = info.get("speed_mean", 0.0) | |
| roll = info.get("rolling_pass_frac", 0.0) | |
| return "stop", (f"RIGOROUS VERDICT WALKS: frame_alt={fa} L={l_d} " | |
| f"R={r_d} spd={sp} roll={roll} -> genuine human gait") | |
| return "continue", (f"RIGOROUS VERDICT NO-WALK " | |
| f"(verdict={verdict}) -> keep training") | |
| def propose_change(state, info): | |
| """Turn a NO-WALK verdict into ONE whitelisted code change. | |
| The reflector (disabled here) would read the verdict + video and return a | |
| structured change dict. When disabled we use a CAUSAL fallback: inspect the | |
| measured verdict and pick the change that addresses the actual failure | |
| mode, falling back to a safe rotation over the whitelist. Every applied | |
| change is git-committed + test-gated inside loop_changes.apply_change(); a | |
| bad change auto-reverts and we just retrain identical weights. | |
| """ | |
| info = info or {} | |
| l_d = info.get("left_duty_mean", 0.0) | |
| r_d = info.get("right_duty_mean", 0.0) | |
| c_alt = info.get("contact_alt_mean", 0.0) | |
| switch = info.get("switch_hz_mean", 0.0) | |
| # One-leg dominance / dead leg: the failing foot barely contacts while the | |
| # other is always down (high contact asymmetry). Break it with duty_balance. | |
| min_duty = min(l_d, r_d) | |
| if min_duty < 0.2 or c_alt > 0.5: | |
| return {"target": "reward_weight", "term": "duty_balance", "weight": 3.0} | |
| # Hopping came back (high switch_hz but no real double-support / duty fails): | |
| # strengthen the anti-hop term. | |
| if switch > 2.0: | |
| return {"target": "reward_weight", "term": "both_feet_air_time", | |
| "weight": 6.0} | |
| # SQUAT-FREEZE (round-2 failure): both feet planted, no stance transfer, | |
| # no double-support, tiny speed, deep knee bend. The planted double-stance | |
| # earns neither no_fly (needs exactly-one-down) nor feet_moving (needs any | |
| # lift), so raise both to push the policy into alternating single support. | |
| dsup = info.get("double_support_frac", 1.0) | |
| knee = info.get("knee_flex_min_rad", 0.0) | |
| spd = info.get("speed_mean", 0.0) | |
| if switch < 0.4 and dsup < 0.1 and spd < 0.1 and knee > 0.5: | |
| return {"target": "reward_weight", "term": "no_fly", "weight": 4.0} | |
| if switch < 0.4 and dsup < 0.1: | |
| return {"target": "reward_weight", "term": "feet_moving", "weight": 3.0} | |
| # Otherwise rotate over safe structural nudges. | |
| if not getattr(propose_change, "_rot", None): | |
| propose_change._rot = iter([ | |
| {"target": "reward_weight", "term": "air_time", "weight": 4.0}, | |
| {"target": "reward_weight", "term": "foot_clearance", "weight": -0.3}, | |
| {"target": "symmetry", "field": "mirror_loss_coeff", "value": 0.8}, | |
| {"target": "reward_weight", "term": "upright", "weight": 3.0}, | |
| {"target": "reward_weight", "term": "duty_balance", "weight": 3.0}, | |
| {"target": "env_param", "param": "fell_over_limit_angle", "value": 35.0}, | |
| ]) | |
| try: | |
| change = next(propose_change._rot) | |
| except StopIteration: | |
| propose_change._rot = None | |
| return propose_change(state, info) # restart rotation | |
| return change | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--from-run", default=None) | |
| ap.add_argument("--from-checkpoint", default=None) | |
| ap.add_argument("--rounds", type=int, default=0, | |
| help="max rounds; 0 = infinite loop until WALKS (default)") | |
| ap.add_argument("--reflect-only", action="store_true") | |
| ap.add_argument("--no-reflect", action="store_true", | |
| help="skip the opencode reflector; always use the fallback " | |
| "strategy (pure resume, no blocking on muse)") | |
| args = ap.parse_args() | |
| state = load_state() | |
| run = args.from_run or state["run"] | |
| ckpt = args.from_checkpoint or state["checkpoint"] | |
| if not run or not ckpt: | |
| if state.get("rounds_done", 0) == 0 and state.get("run") is None: | |
| run, ckpt = None, None # first-ever round trains fresh from init | |
| else: | |
| print("need --from-run and --from-checkpoint (or a saved state)") | |
| sys.exit(2) | |
| if args.reflect_only: | |
| st = reflect(state) | |
| print("REFLECT ->", json.dumps(st)) | |
| sys.exit(0) | |
| def pick_strategy(state): | |
| if args.no_reflect: | |
| # Reflector disabled: use the safe fallback (continue the same | |
| # microduck-exact config, pure resume). Never blocks on muse. | |
| return {"name": f"continue-{state['rounds_done'] + 1}", | |
| "kind": "resume-train", "iters": 1000, "envs": 4096, | |
| "why": "reflector disabled; continuing microduck-exact resume", | |
| "change": "none (no reflector)"} | |
| if state.get("rounds_done", 0) == 0 and "strategy" not in state: | |
| return SEED_STRATEGY | |
| return reflect(state) | |
| for step in range(args.rounds) if args.rounds and args.rounds > 0 else iter(int, 1): | |
| st = pick_strategy(state) | |
| append_strategy(st) | |
| fresh = (state.get("run") is None) # first ever run = fresh from init | |
| log_row("loop", f"round {state['rounds_done'] + 1}: {st['name']} " | |
| f"{'FRESH' if fresh else 'from ' + ckpt}", | |
| st.get("why", ""), f"logs/rsl_rl/velocity/{run}/{ckpt}", | |
| "training") | |
| new_run = train_round(st, run, ckpt, fresh=fresh) | |
| if new_run is None: | |
| log_row("loop", f"{st['name']} failed to produce a run", "see loop log", | |
| f"/tmp/opencode/loop_{st['name']}.log", "failed") | |
| sys.exit(2) | |
| new_ckpt = newest_model(new_run) | |
| if new_ckpt is None: | |
| log_row("loop", f"{st['name']} produced no checkpoint (training died)", | |
| "see loop log", f"/tmp/opencode/loop_{st['name']}.log", "failed") | |
| sys.exit(2) | |
| try: | |
| passed, info = verify(new_run, new_ckpt) | |
| except RuntimeError as e: | |
| log_row("loop", f"verify failed: {e}", "training died early", | |
| f"logs/rsl_rl/velocity/{new_run}", "failed") | |
| sys.exit(2) | |
| # ANY run we launched is now recorded so it can always be resumed. | |
| state.update({"rounds_done": state["rounds_done"] + 1, | |
| "run": new_run, "checkpoint": new_ckpt, | |
| "last_verdict": info, "last_video": info.get("video", "")}) | |
| save_state(state) | |
| run, ckpt = new_run, new_ckpt | |
| # Agentic loop-condition: subagent decides STOP (walks) vs CONTINUE. | |
| decision, why = decide(state, new_run, new_ckpt) | |
| if decision == "stop": | |
| print(f"LOOP VERDICT: WALKS at {new_run}/{new_ckpt} ({why})") | |
| log_row("loop", f"WALKS at {new_ckpt}", | |
| why, f"logs/rsl_rl/velocity/{new_run}/{new_ckpt}", "done") | |
| sys.exit(0) | |
| print(f"round {state['rounds_done']} NO-WALK -> {why} -> continuing (infinite loop)") | |
| # Agentic code-change: propose ONE whitelisted edit from the verdict and | |
| # apply it (git-committed + test-gated; auto-reverts on failure) so the | |
| # next round trains on a nudged config instead of identical weights. | |
| change = propose_change(state, info) | |
| res = loop_changes.apply_change(change) | |
| print(f" code-change: {res}") | |
| # brief pause so we never spin on a bad state; then loop forever. | |
| time.sleep(5) | |
| # Round budget exhausted without a WALKS verdict -> continue (not a win). | |
| # The supervisor must NOT read this as WALKS. Exit 1 = keep looping. | |
| print(f"round budget spent ({state['rounds_done']} rounds) -> NO-WALK, continuing") | |
| sys.exit(1) | |
| if __name__ == "__main__": | |
| main() | |