Initial commit
This commit is contained in:
@@ -0,0 +1 @@
|
||||
"""S10 policies, world models, replay, imagined rollouts and control UI."""
|
||||
@@ -0,0 +1,88 @@
|
||||
"""Explicit name-based mapping between mjlab and canonical S10 model tensors."""
|
||||
|
||||
import torch
|
||||
|
||||
from s10_supermimic.robot.config import (
|
||||
LEG_JOINT_NAMES, WHEEL_BODY_NAMES, WHEEL_GEOM_NAMES, WHEEL_JOINT_NAMES,
|
||||
)
|
||||
from .world_model import STATE_DIM, ACTION_DIM
|
||||
|
||||
|
||||
class S10Bridge:
|
||||
def __init__(self, env):
|
||||
self.env = env
|
||||
self.robot = env.scene["robot"]
|
||||
self.legs, _ = self.robot.find_joints(LEG_JOINT_NAMES, preserve_order=True)
|
||||
self.wheels, _ = self.robot.find_joints(WHEEL_JOINT_NAMES, preserve_order=True)
|
||||
if len(self.legs) != 12 or len(self.wheels) != 4:
|
||||
raise ValueError("S10 requires 12 leg joints and 4 wheel joints")
|
||||
pos = env.action_manager.get_term("joint_pos")
|
||||
vel = env.action_manager.get_term("wheel_vel")
|
||||
if list(env.cfg.actions) != ["joint_pos", "wheel_vel"]:
|
||||
raise ValueError("Expected joint_pos then wheel_vel action terms")
|
||||
names = list(pos.target_names) + list(vel.target_names)
|
||||
canonical = LEG_JOINT_NAMES + WHEEL_JOINT_NAMES
|
||||
if len(names) != ACTION_DIM or set(names) != set(canonical):
|
||||
raise ValueError("Unexpected S10 action targets")
|
||||
self.to_canonical = [names.index(n) for n in canonical]
|
||||
self.to_env = [canonical.index(n) for n in names]
|
||||
geoms, _ = self.robot.find_geoms(WHEEL_GEOM_NAMES, preserve_order=True)
|
||||
self.geoms = self.robot.indexing.geom_ids[geoms]
|
||||
body, _ = self.robot.find_bodies("base_link")
|
||||
self.base = self.robot.indexing.body_ids[body[0]]
|
||||
# XML reference, not an already-randomized simulator value after reset.
|
||||
from s10_supermimic.robot.config import get_spec
|
||||
self.nominal_com = torch.as_tensor(get_spec().body("base_link").ipos.copy(),
|
||||
device=env.device, dtype=torch.float32)
|
||||
|
||||
def canonical_action(self, action):
|
||||
return action[:, self.to_canonical]
|
||||
|
||||
def env_action(self, action):
|
||||
return action[:, self.to_env]
|
||||
|
||||
def state(self):
|
||||
d = self.robot.data
|
||||
s = torch.cat((d.root_link_pos_w[:, 2:3], d.projected_gravity_b,
|
||||
d.root_link_lin_vel_b, d.root_link_ang_vel_b,
|
||||
(d.joint_pos - d.default_joint_pos)[:, self.legs],
|
||||
d.joint_vel[:, self.legs], d.joint_vel[:, self.wheels]), -1)
|
||||
assert s.shape[-1] == STATE_DIM
|
||||
return s
|
||||
|
||||
def pose(self):
|
||||
d = self.robot.data
|
||||
return torch.cat((d.root_link_pos_w[:, :2], d.heading_w[:, None]), -1)
|
||||
|
||||
def context(self):
|
||||
model = self.env.sim.model
|
||||
friction = model.geom_friction
|
||||
com = model.body_ipos
|
||||
mu = friction[:, self.geoms, 0] if friction.ndim == 3 else friction[self.geoms, 0].expand(self.env.num_envs, -1)
|
||||
offset = com[:, self.base] if com.ndim == 3 else com[self.base].expand(self.env.num_envs, -1)
|
||||
return torch.cat((mu, offset - self.nominal_com), -1)
|
||||
|
||||
def contacts(self):
|
||||
"""Endpoint targets: contact flags4, net wheel forces12 in base frame (N).
|
||||
|
||||
Explicit names prevent silent FL/FR swaps in sensor expansion order.
|
||||
Terminal/reset endpoints are excluded by replay validity in the learner.
|
||||
"""
|
||||
from mjlab.utils.lab_api.math import quat_apply_inverse
|
||||
sensor = self.env.scene["wheel_ground_contact"]
|
||||
names = [name.split('/')[-1] for name in sensor.primary_names]
|
||||
order = [names.index(name) for name in WHEEL_BODY_NAMES]
|
||||
found = (sensor.data.found[:,order] > 0).float().reshape(-1,4)
|
||||
force = sensor.data.force[:,order].reshape(-1,4,3)
|
||||
quat = self.robot.data.root_link_quat_w[:,None,:].expand(-1,4,-1)
|
||||
force = quat_apply_inverse(quat.reshape(-1,4),force.reshape(-1,3)).reshape(-1,12)
|
||||
return torch.cat((found,force),-1)
|
||||
|
||||
|
||||
def pose_delta(before, after):
|
||||
dxy = after[:, :2] - before[:, :2]
|
||||
c, s = before[:, 2].cos(), before[:, 2].sin()
|
||||
yaw = after[:, 2] - before[:, 2]
|
||||
return torch.stack((c * dxy[:, 0] + s * dxy[:, 1],
|
||||
-s * dxy[:, 0] + c * dxy[:, 1],
|
||||
torch.atan2(yaw.sin(), yaw.cos())), -1)
|
||||
@@ -0,0 +1,24 @@
|
||||
"""Non-wheel collision termination and pre-reset labels for imagined survival."""
|
||||
import torch
|
||||
|
||||
|
||||
def illegal_nonwheel_contact(env, force_threshold=10.):
|
||||
data = env.scene['nonwheel_ground_contact'].data
|
||||
force = data.force_history if data.force_history is not None else data.force
|
||||
finite = torch.isfinite(force).flatten(1).all(-1)
|
||||
magnitude = torch.nan_to_num(force, nan=0., posinf=0., neginf=0.).norm(dim=-1)
|
||||
peak = magnitude.flatten(1).amax(-1)
|
||||
hit = (peak > force_threshold) & finite
|
||||
# Scene/termination buffers are cleared by automatic reset. These copies
|
||||
# survive it, so the learner retains the transition that caused termination.
|
||||
env._nonwheel_hit = hit.detach().clone()
|
||||
env._nonwheel_label_valid = finite.detach().clone()
|
||||
env._nonwheel_peak = peak.detach().clone()
|
||||
return hit
|
||||
|
||||
|
||||
def collision_events(batch):
|
||||
valid = batch['valid'].squeeze(-1).bool()
|
||||
prefix = torch.cat((torch.ones_like(valid[:,:1]), valid.cumprod(1).bool()[:,:-1]),1)
|
||||
return ((batch['collision'].squeeze(-1) > .5)
|
||||
& batch['collision_valid'].squeeze(-1).bool() & prefix).any(-1)
|
||||
@@ -0,0 +1,32 @@
|
||||
"""Velocity range presets. V4 is time-based; v5/v6/v7 are stair-gated.
|
||||
|
||||
S10 v7 uses ``V7_STAGES`` as the 1.5/0.8/1.2 cap, but stage advancement comes
|
||||
from ``stair_curriculum.next_stage`` (actual up-stair pass rate). Only v4
|
||||
advances by iteration count alone.
|
||||
"""
|
||||
STAGES = (
|
||||
(0, (-.3,.8), (-.2,.2), (-.8,.8)),
|
||||
(1000, (-.5,1.), (-.4,.4), (-1.,1.)),
|
||||
(2500, (-.8,1.5), (-.6,.6), (-1.5,1.5)),
|
||||
(5000, (-1.5,1.8), (-.8,.8), (-2.,2.)),
|
||||
(7500, (-2.,2.), (-1.,1.), (-3.,3.)),
|
||||
)
|
||||
V7_STAGES = (
|
||||
(0, (-.3,.8), (-.2,.2), (-.5,.5)),
|
||||
(1000, (-.5,1.), (-.4,.4), (-.7,.7)),
|
||||
(2500, (-.8,1.2), (-.6,.6), (-.9,.9)),
|
||||
(5000, (-1.2,1.4), (-.7,.7), (-1.1,1.1)),
|
||||
(7500, (-1.5,1.5), (-.8,.8), (-1.2,1.2)),
|
||||
)
|
||||
|
||||
|
||||
def stages_for(revision):
|
||||
return V7_STAGES if revision >= 7 else STAGES
|
||||
|
||||
|
||||
def set_command_stage(cfg, iteration, stage=None, stages=STAGES):
|
||||
index = stage if stage is not None else max(
|
||||
i for i, item in enumerate(stages) if max(0, iteration) >= item[0])
|
||||
_, x, y, yaw = stages[index]
|
||||
cfg.ranges.lin_vel_x, cfg.ranges.lin_vel_y, cfg.ranges.ang_vel_z = x, y, yaw
|
||||
return index
|
||||
@@ -0,0 +1,114 @@
|
||||
<!doctype html>
|
||||
<html lang="zh-CN">
|
||||
<head>
|
||||
<meta charset="utf-8"><meta name="viewport" content="width=device-width,initial-scale=1">
|
||||
<title>S10 · 运行控制</title>
|
||||
<style>
|
||||
:root{color-scheme:dark;font-family:system-ui,-apple-system,sans-serif;background:#10141b;color:#e4eaf3}*{box-sizing:border-box}body{margin:0;padding:28px}header{display:flex;justify-content:space-between;align-items:center;margin-bottom:24px}h1{font-size:25px;margin:0}header small{color:#93a1b5;font-size:13px}.status{background:#202a36;padding:8px 14px;border-radius:20px;color:#a2b7cb;font-size:13px}main{display:grid;grid-template-columns:minmax(400px,1fr) 350px;gap:22px;max-width:1600px;margin:auto}.card{background:#19212d;border:1px solid #2c3748;border-radius:16px;overflow:hidden}.viewer{background:#080b10;display:flex;align-items:center;justify-content:center;min-height:400px}.viewer img{display:block;width:100%;aspect-ratio:4/3;object-fit:contain}.caption{padding:16px;color:#9cabc0;font-size:13px}.panel{padding:24px}h2{font-size:18px;margin:0 0 8px}.muted{color:#9cabc0;font-size:13px;line-height:1.7}.slider{margin:26px 0}.slider label{display:flex;justify-content:space-between;align-items:center;font-size:14px}output{font:600 21px ui-monospace,monospace;color:#64debf}input[type=range]{width:100%;accent-color:#58cfb1;margin:17px 0 5px;cursor:pointer}.bounds{display:flex;justify-content:space-between;color:#8292a7;font-size:12px}.actions{display:flex;gap:10px;margin:12px 0}button{cursor:pointer;border:1px solid #415067;border-radius:9px;background:#273448;color:#edf3fa;padding:11px 14px;font:inherit;font-size:14px;flex:1}button:hover{filter:brightness(1.2)}button.stop{background:#94404c;border-color:#c05b69;width:100%;font-weight:600}.telemetry{border-top:1px solid #344054;padding-top:20px;margin-top:24px}.metric{display:flex;justify-content:space-between;padding:8px 0;font-size:14px}.metric b{font-family:ui-monospace,monospace}#wm{font-size:11px;max-width:210px;text-align:right;white-space:normal;line-height:1.35}.dot{color:#64debf}.note{margin-top:20px;font-size:12px;color:#8a9ab0;line-height:1.7}@media(max-width:900px){body{padding:16px}main{grid-template-columns:1fr}.viewer{min-height:200px}}
|
||||
select{background:#273448;color:#edf3fa;padding:8px;border:1px solid #415067;border-radius:6px;margin:5px}.viewer img{cursor:grab;touch-action:none;user-select:none}.viewer img:active{cursor:grabbing}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<header><div><h1>S10 <span style="font-weight:400;color:#9cabc0">/ 运行控制</span></h1><small id="model">正在加载策略…</small></div><div class="status" id="connection">连接中</div></header>
|
||||
<main>
|
||||
<section class="card"><div class="viewer"><img id="frame" src="/frame.jpg" alt="S10 实时仿真画面"></div><div class="caption" id="caption">单环境 · CPU · 实时画面</div></section>
|
||||
<section class="card panel">
|
||||
<h2>地形与视角</h2>
|
||||
<div><select id="terrainKind" aria-label="地形"><option value="flat">平地</option><option value="up">上楼梯</option><option value="down">下楼梯</option><option value="platform">障碍平台</option><option value="slope">斜坡</option><option value="rubble">碎石地</option></select>
|
||||
<select id="terrainLevel" aria-label="难度"><option value="0">等级 0 · 2 cm / 3°</option><option value="1">等级 1 · 4 cm / 5.4°</option><option value="2">等级 2 · 8 cm / 7.8°</option><option value="3">等级 3 · 12 cm / 10.2°</option><option value="4">等级 4 · 16 cm / 12.6°</option><option value="5">等级 5 · 20 cm / 15°</option></select></div>
|
||||
<div class="actions"><button id="changeScene">切换地形</button><button id="resetCamera">恢复视角</button></div>
|
||||
<div class="muted">在画面上按住左键拖动旋转，滚轮缩放，双击恢复视角。镜头始终跟随机器人。切换地形会重置机器人。</div>
|
||||
|
||||
<h2>键盘控制 · 速度大小</h2><div class="muted">↑↓ 前进/后退 · ←→ 横移 · Q/E 转向<br>按住移动，松开停止；滑条只调速度大小。</div><div class="muted" id="requested">当前指令：0 / 0 / 0</div>
|
||||
<div class="slider"><label for="vx">前进 / 后退 · vx <output id="vxout">0.00</output></label><input id="vx" type="range" min="0" max="2" step="0.01" value="0.3"><div class="bounds"><span>0 m/s</span><span>+2.00 m/s</span></div></div>
|
||||
<div class="slider"><label for="vy">左移 / 右移 · vy <output id="vyout">0.00</output></label><input id="vy" type="range" min="0" max="1" step="0.01" value="0.2"><div class="bounds"><span>0 m/s</span><span>+1.00 m/s</span></div></div>
|
||||
<div class="slider"><label for="yaw">左转 / 右转 · yaw rate <output id="yawout">0.00</output></label><input id="yaw" type="range" min="0" max="3" step="0.01" value="0.4"><div class="bounds"><span>0 rad/s</span><span>+3.00 rad/s</span></div></div>
|
||||
<button class="stop" id="stop">停止 · 空格</button><div class="actions"><button id="pause">暂停仿真</button><button id="reset">重置机器人</button></div>
|
||||
<div class="telemetry"><h2>实际速度</h2><div class="metric"><span>vx</span><b id="actual0">—</b></div><div class="metric"><span>vy</span><b id="actual1">—</b></div><div class="metric"><span>yaw rate</span><b id="actual2">—</b></div><div class="metric"><span>WM 预测误差</span><b id="wm">—</b></div></div>
|
||||
<div class="note" id="note">黄→橙→红渐变残影是世界模型从当前状态想象的未来 32 步（默认每 4 步画一个）；WM 预测误差会随仿真推进与真实状态对齐后更新。切换标签页或关闭页面会归零。连接中断超过 2 秒也会归零。暂停会清空指令。</div>
|
||||
</section>
|
||||
</main>
|
||||
<script>
|
||||
const sliders=['vx','vy','yaw'].map(id=>document.getElementById(id));
|
||||
const connection=document.getElementById('connection');
|
||||
let dirty=false,busy=false,actionQueue=[],dragging=false;
|
||||
let cameraDelta=[0,0,0],pointer=null,sceneInitialized=false;
|
||||
const keys=new Set(),motionKeys=new Set(["ArrowUp","ArrowDown","ArrowLeft","ArrowRight","KeyQ","KeyE"]);
|
||||
function command(){
|
||||
const x=Number(keys.has("ArrowUp"))-Number(keys.has("ArrowDown"));
|
||||
const y=Number(keys.has("ArrowLeft"))-Number(keys.has("ArrowRight"));
|
||||
const z=Number(keys.has("KeyQ"))-Number(keys.has("KeyE"));
|
||||
const scale=Math.max(1,Math.hypot(x,y));
|
||||
return [Number(sliders[0].value)*x/scale,Number(sliders[1].value)*y/scale,Number(sliders[2].value)*z];
|
||||
}
|
||||
function labels(){sliders.forEach(s=>document.getElementById(s.id+'out').textContent=Number(s.value).toFixed(2))}
|
||||
function zero(){keys.clear();dirty=true}
|
||||
function action(name){if(name!=='heartbeat'){zero();actionQueue.push(name)}}
|
||||
sliders.forEach(s=>{s.addEventListener('input',()=>{labels();dirty=true});s.addEventListener('pointerdown',()=>dragging=true)});
|
||||
window.addEventListener('pointerup',()=>dragging=false);
|
||||
document.getElementById('stop').onclick=()=>action('stop');
|
||||
document.getElementById('pause').onclick=()=>action('pause');
|
||||
document.getElementById('reset').onclick=()=>action('reset');
|
||||
window.addEventListener('keydown',e=>{
|
||||
if(e.code==='Space'){e.preventDefault();action('stop');return}
|
||||
if(motionKeys.has(e.code)){e.preventDefault();if(!e.repeat){keys.add(e.code);dirty=true}}
|
||||
});
|
||||
window.addEventListener('keyup',e=>{
|
||||
if(motionKeys.has(e.code)){e.preventDefault();keys.delete(e.code);dirty=true}
|
||||
});
|
||||
window.addEventListener('blur',()=>{endDrag();action('stop');leave()});
|
||||
labels();
|
||||
function leave(){zero();navigator.sendBeacon('/api/control',JSON.stringify({action:'stop'}))}
|
||||
document.addEventListener('visibilitychange',()=>{if(document.hidden)leave()});
|
||||
window.addEventListener('pagehide',leave);
|
||||
function show(s){
|
||||
connection.textContent=s.loading?'正在切换地形…':s.ready?(s.paused?'已暂停':'● 已连接'):'仿真初始化中';
|
||||
if(!sceneInitialized && s.terrain_kind){
|
||||
document.getElementById('terrainKind').value=s.terrain_kind;
|
||||
document.getElementById('terrainLevel').value=String(s.terrain_level);sceneInitialized=true;
|
||||
}
|
||||
document.getElementById('changeScene').disabled=!!s.loading||!s.terrain;
|
||||
document.getElementById('pause').textContent=s.paused?'继续仿真':'暂停仿真';
|
||||
document.getElementById('requested').textContent='当前指令：vx '+s.command[0].toFixed(2)+' / vy '+s.command[1].toFixed(2)+' / yaw '+s.command[2].toFixed(2);
|
||||
(s.actual||[]).forEach((v,i)=>document.getElementById('actual'+i).textContent=Number(v).toFixed(3)+(i===2?' rad/s':' m/s'));
|
||||
document.getElementById('wm').textContent=s.wm_error||'—';
|
||||
if(s.checkpoint)document.getElementById('model').textContent=s.checkpoint+' · 训练步数 '+s.iteration;
|
||||
if(s.scene)document.getElementById('caption').textContent=s.scene+' · 单环境 · '+s.device+' · 渲染 '+Number(s.render_fps||0).toFixed(1)+' FPS · 实时倍率 '+Number(s.realtime_factor||0).toFixed(2)+'×';
|
||||
if(s.terrain)document.getElementById('note').textContent=(s.terrain_revision>=3?'新版训练指令范围：vx ±2 m/s，vy ±1 m/s，yaw ±3 rad/s。实际能力以评估为准。':'当前为前向地形策略。训练范围：vx 0.2–0.45 m/s，vy 0，yaw ±0.05 rad/s；反向和大幅转向尚未验证。')+' 黄→橙→红渐变残影是世界模型预测的未来 32 步；WM 误差在与真实状态对齐后更新。切换标签页或断开连接后指令归零。';
|
||||
}
|
||||
async function tick(){
|
||||
if(busy||document.hidden)return;busy=true;
|
||||
try{
|
||||
let body={action:'heartbeat'};
|
||||
if(actionQueue.length){const item=actionQueue.shift();body=typeof item==='string'?{action:item}:item;}
|
||||
else if(dirty){body={action:'command',values:command()};dirty=false}
|
||||
if(cameraDelta.some(v=>v!==0)){body.camera_delta=cameraDelta;cameraDelta=[0,0,0]}
|
||||
const response=await fetch('/api/control',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify(body),signal:AbortSignal.timeout(1500)});
|
||||
if(!response.ok)throw Error(response.status);show(await response.json());
|
||||
}catch(e){connection.textContent='连接已断开';zero()}finally{busy=false}
|
||||
}
|
||||
setInterval(tick,33);tick();
|
||||
const image=document.getElementById('frame');
|
||||
image.draggable=false;
|
||||
image.addEventListener('pointerdown',e=>{
|
||||
if(e.button!==0)return;e.preventDefault();pointer={id:e.pointerId,x:e.clientX,y:e.clientY};image.setPointerCapture(e.pointerId);
|
||||
});
|
||||
image.addEventListener('pointermove',e=>{
|
||||
if(!pointer||pointer.id!==e.pointerId)return;
|
||||
cameraDelta[0]-=(e.clientX-pointer.x)*.4;cameraDelta[1]-=(e.clientY-pointer.y)*.3;
|
||||
pointer.x=e.clientX;pointer.y=e.clientY;
|
||||
});
|
||||
function endDrag(){pointer=null}
|
||||
image.addEventListener('pointerup',endDrag);image.addEventListener('pointercancel',endDrag);image.addEventListener('lostpointercapture',endDrag);
|
||||
|
||||
image.addEventListener('wheel',e=>{e.preventDefault();cameraDelta[2]+=Math.max(-1,Math.min(1,e.deltaY*(e.deltaMode===1?.03:e.deltaMode===2?.5:.001)));},{passive:false});
|
||||
function resetCamera(){cameraDelta=[0,0,0];actionQueue.push('camera_reset')}
|
||||
image.addEventListener('dblclick',resetCamera);
|
||||
document.getElementById('resetCamera').onclick=resetCamera;
|
||||
document.getElementById('changeScene').onclick=()=>{
|
||||
zero();actionQueue.push({action:'scene',kind:document.getElementById('terrainKind').value,level:Number(document.getElementById('terrainLevel').value)});
|
||||
};
|
||||
function nextFrame(){setTimeout(()=>{if(!document.hidden)image.src='/frame.jpg?t='+Date.now();else nextFrame()},16)}
|
||||
image.onload=nextFrame;image.onerror=nextFrame;if(image.complete)nextFrame();
|
||||
</script>
|
||||
</body></html>
|
||||
@@ -0,0 +1,171 @@
|
||||
"""Local browser controls. Latest-frame storage only; simulation owns MuJoCo."""
|
||||
import errno
|
||||
import json
|
||||
import math
|
||||
from pathlib import Path
|
||||
import threading
|
||||
import time
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
|
||||
|
||||
class ControlState:
|
||||
def __init__(self, timeout=2.):
|
||||
self.speed, self.lateral_speed, self.turn_speed = 2., 1., 3.
|
||||
self.lock = threading.RLock()
|
||||
self.values = (0., 0., 0.)
|
||||
self.last_seen = time.monotonic()
|
||||
self.timeout = timeout
|
||||
self.paused = False
|
||||
self.reset_requested = False
|
||||
self.camera = [140., -20., 2.5]
|
||||
self.scene_requested = None
|
||||
self.loading = False
|
||||
self.jpeg = b""
|
||||
self.status = {"ready": False, "actual": [0., 0., 0.]}
|
||||
|
||||
def clear(self):
|
||||
with self.lock:
|
||||
self.values = (0., 0., 0.)
|
||||
|
||||
def command(self):
|
||||
with self.lock:
|
||||
if time.monotonic() - self.last_seen > self.timeout:
|
||||
self.clear()
|
||||
return (0., 0., 0.) if self.paused or self.loading else self.values
|
||||
|
||||
def update(self, request):
|
||||
if not isinstance(request, dict):
|
||||
raise ValueError("Expected an object")
|
||||
action = request.get("action", "command")
|
||||
if action not in ("command", "heartbeat", "stop", "pause", "reset", "camera_reset", "scene"):
|
||||
raise ValueError("Unknown action")
|
||||
values = None
|
||||
if action == "command":
|
||||
values = request.get("values")
|
||||
if not isinstance(values, list) or len(values) != 3:
|
||||
raise ValueError("Expected vx, vy, yaw_rate")
|
||||
if any(isinstance(v, bool) or not isinstance(v, (int, float)) or not math.isfinite(v) for v in values):
|
||||
raise ValueError("Commands must be finite numbers")
|
||||
values = tuple(max(-limit, min(limit, value)) for value, limit in
|
||||
zip(values, (self.speed, self.lateral_speed, self.turn_speed)))
|
||||
camera_delta = request.get("camera_delta", [0., 0., 0.])
|
||||
if (not isinstance(camera_delta, list) or len(camera_delta) != 3
|
||||
or any(isinstance(v, bool) or not isinstance(v, (int, float)) or not math.isfinite(v) for v in camera_delta)):
|
||||
raise ValueError("Invalid camera delta")
|
||||
scene = None
|
||||
if action == "scene":
|
||||
kind, level = request.get("kind"), request.get("level")
|
||||
if kind not in ("flat", "up", "down", "platform", "slope", "rubble") or type(level) is not int or level not in range(6):
|
||||
raise ValueError("Invalid terrain selection")
|
||||
if not self.status.get("terrain", False):
|
||||
raise ValueError("Terrain selection requires a terrain policy")
|
||||
scene = (kind, level)
|
||||
with self.lock:
|
||||
self.camera[0] = (self.camera[0] + camera_delta[0]) % 360
|
||||
self.camera[1] = max(-89., min(0., self.camera[1] + camera_delta[1]))
|
||||
self.camera[2] = max(.5, min(12., self.camera[2] * math.exp(max(-2., min(2., camera_delta[2])))))
|
||||
if action == "camera_reset":
|
||||
self.camera = [140., -20., 2.5]
|
||||
if scene is not None:
|
||||
self.scene_requested = scene
|
||||
self.loading = True
|
||||
self.clear()
|
||||
# Expire stale commands before a late heartbeat can revive them.
|
||||
self.command()
|
||||
self.last_seen = time.monotonic()
|
||||
if values is not None and not self.loading:
|
||||
self.values = values
|
||||
elif action in ("stop", "reset", "pause"):
|
||||
self.clear()
|
||||
if action == "pause":
|
||||
self.paused = not self.paused
|
||||
elif action == "reset":
|
||||
self.reset_requested = True
|
||||
return self.snapshot()
|
||||
|
||||
def snapshot(self):
|
||||
with self.lock:
|
||||
return dict(self.status, command=list(self.command()), paused=self.paused, loading=self.loading, camera=list(self.camera))
|
||||
|
||||
def take_scene(self):
|
||||
with self.lock:
|
||||
scene, self.scene_requested = self.scene_requested, None
|
||||
return scene
|
||||
|
||||
def finish_scene(self):
|
||||
with self.lock:
|
||||
self.loading = self.scene_requested is not None
|
||||
self.clear()
|
||||
|
||||
def take_reset(self):
|
||||
with self.lock:
|
||||
reset, self.reset_requested = self.reset_requested, False
|
||||
return reset
|
||||
|
||||
def publish(self, actual, step, **metadata):
|
||||
with self.lock:
|
||||
self.status.update(ready=True, actual=actual, step=step, **metadata)
|
||||
|
||||
|
||||
def make_server(state, port=8765):
|
||||
html = Path(__file__).with_name("gui.html").read_bytes()
|
||||
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
def log_message(self, *args):
|
||||
pass
|
||||
|
||||
def send(self, status, data, kind):
|
||||
self.send_response(status)
|
||||
self.send_header("Content-Type", kind)
|
||||
self.send_header("Content-Length", str(len(data)))
|
||||
self.send_header("Cache-Control", "no-store")
|
||||
self.end_headers()
|
||||
try:
|
||||
self.wfile.write(data)
|
||||
except (BrokenPipeError, ConnectionResetError):
|
||||
pass
|
||||
|
||||
def do_GET(self):
|
||||
route = self.path.split("?", 1)[0]
|
||||
if route == "/":
|
||||
self.send(200, html, "text/html; charset=utf-8")
|
||||
elif route == "/api/state":
|
||||
self.send(200, json.dumps(state.snapshot()).encode(), "application/json")
|
||||
elif route == "/frame.jpg":
|
||||
with state.lock:
|
||||
jpeg = state.jpeg
|
||||
self.send(200 if jpeg else 503, jpeg, "image/jpeg")
|
||||
else:
|
||||
self.send(404, b"Not found", "text/plain")
|
||||
|
||||
def do_POST(self):
|
||||
if self.path != "/api/control":
|
||||
self.send(404, b"Not found", "text/plain")
|
||||
return
|
||||
# Local-only API: browser requests must come from this page's origin.
|
||||
origin = self.headers.get("Origin")
|
||||
if origin and origin != "http://" + self.headers.get("Host", ""):
|
||||
self.send(403, b"Wrong origin", "text/plain")
|
||||
return
|
||||
try:
|
||||
length = int(self.headers.get("Content-Length", "0"))
|
||||
if not 0 < length <= 2048:
|
||||
raise ValueError("Invalid body length")
|
||||
result = state.update(json.loads(self.rfile.read(length)))
|
||||
self.send(200, json.dumps(result).encode(), "application/json")
|
||||
except (ValueError, TypeError) as error:
|
||||
self.send(400, json.dumps({"error": str(error)}).encode(), "application/json")
|
||||
|
||||
requested_port = port
|
||||
while True:
|
||||
try:
|
||||
server = ThreadingHTTPServer(("127.0.0.1", port), Handler)
|
||||
break
|
||||
except OSError as error:
|
||||
if error.errno != errno.EADDRINUSE or port >= 65535:
|
||||
raise
|
||||
port += 1
|
||||
if port != requested_port:
|
||||
print(f"Port {requested_port} is occupied; using http://127.0.0.1:{server.server_port}/", flush=True)
|
||||
server.daemon_threads = True
|
||||
return server
|
||||
@@ -0,0 +1,198 @@
|
||||
"""Differentiable rewards, full-return actor objective and TD(lambda) critic targets."""
|
||||
|
||||
import math
|
||||
from contextlib import contextmanager
|
||||
import torch
|
||||
from .world_model import GRAVITY, LIN_VEL, ANG_VEL, LEG_POS, LEG_VEL, WHEEL_VEL
|
||||
from .policy import frame, push_history, critic_input
|
||||
|
||||
|
||||
@contextmanager
|
||||
def frozen(*modules):
|
||||
"""Freeze weights without disabling gradients through inputs/actions."""
|
||||
parameters = [(p, p.requires_grad) for m in modules for p in m.parameters()]
|
||||
try:
|
||||
for p, _ in parameters:
|
||||
p.requires_grad_(False)
|
||||
yield
|
||||
finally:
|
||||
for p, enabled in parameters:
|
||||
p.requires_grad_(enabled)
|
||||
|
||||
|
||||
def survival(state):
|
||||
gravity = torch.nn.functional.normalize(state[:, GRAVITY], dim=-1, eps=1e-6)
|
||||
# Smooth proxy for the environment's 70-degree fall cutoff, width ~5 deg.
|
||||
return torch.sigmoid((-gravity[:, 2] - math.cos(math.radians(70))) / 0.08)
|
||||
|
||||
|
||||
def continuation_value(state, constants):
|
||||
"""Imagined continuation probability for the current task.
|
||||
|
||||
The v4 proxy treats a 70-degree tilt as an immediate fall. The upward task
|
||||
instead terminates only after 4 s of bad orientation, which is 200 control
|
||||
steps, far beyond the 32/64-step imagination horizon, and the counter is not
|
||||
part of the state. Applying the v4 proxy there truncates the return after
|
||||
the first step whenever the robot is inverted, so the actor cannot learn to
|
||||
flip back.
|
||||
"""
|
||||
if constants.get("upright_mode") == "upward":
|
||||
return torch.ones(state.shape[0], dtype=state.dtype, device=state.device)
|
||||
return survival(state)
|
||||
|
||||
|
||||
def reward_terms(state, action, previous, command, constants, survival_prob, scan=None):
|
||||
v, w = state[:, LIN_VEL], state[:, ANG_VEL]
|
||||
q, qd, wd = state[:, LEG_POS], state[:, LEG_VEL], state[:, WHEEL_VEL]
|
||||
standing = (command.square().sum(-1) < 1e-6).float()
|
||||
# Endpoint PD torque surrogate: actual MuJoCo forces also reflect implicit
|
||||
# integration over substeps. The low-weight effort term is approximate.
|
||||
leg_torque = (constants["kp"] * (constants.get("leg_action_scale", 0.25) * action[:, :12] - q)
|
||||
- constants["kd"] * qd).clamp(-constants["effort"], constants["effort"])
|
||||
wheel_limit = constants.get("wheel_velocity_limit", 35.)
|
||||
wheel_effort = constants.get("wheel_effort", 45.)
|
||||
wheel_torque = (constants.get("wheel_kd", 2.) * ((constants.get("wheel_action_scale", 15.) * action[:, 12:]).clamp(-wheel_limit, wheel_limit) - wd)).clamp(-wheel_effort, wheel_effort)
|
||||
absolute_q = q + constants["default_q"]
|
||||
limits = ((constants["limits"][..., 0] - absolute_q).clamp_min(0)
|
||||
+ (absolute_q - constants["limits"][..., 1]).clamp_min(0)).sum(-1)
|
||||
posture = q.square().sum(-1)
|
||||
if constants.get("terrain_posture_v2", False) and not constants.get("wheel_local_rewards", False):
|
||||
if scan is None:
|
||||
raise ValueError("Terrain v2 posture reward requires the privileged scan")
|
||||
from .terrain_posture import posture_cost
|
||||
posture = posture_cost(q, scan, constants.get("posture_relief_floor", 0.125))
|
||||
terrain = constants.get("terrain_rewards", False)
|
||||
gravity = state[:, GRAVITY]
|
||||
if constants.get("upright_mode") == "upward":
|
||||
# Sign-aware recovery reward: 1 upright, 0.5 sideways, 0 upside down.
|
||||
upright = (0.5 * (1. - gravity[:, 2])).clamp(0., 1.)
|
||||
else:
|
||||
upright = torch.exp(
|
||||
-gravity[:, :2].square().sum(-1) / constants.get("upright_std", 0.45)**2)
|
||||
terms = {
|
||||
"track_linear_velocity": torch.exp(-((v[:, :2] - command[:, :2]).square().sum(-1)
|
||||
+ (0. if terrain else v[:, 2].square())) / 0.25),
|
||||
"track_angular_velocity": torch.exp(-((w[:, 2] - command[:, 2]).square()
|
||||
+ (0. if terrain else w[:, :2].square().sum(-1))) / 0.5),
|
||||
"upright": upright,
|
||||
"leg_posture": posture,
|
||||
"standstill": standing * (v.square().sum(-1) + w.square().sum(-1)
|
||||
+ 0.01 * wd.square().sum(-1)),
|
||||
"action_rate_l2": (action - previous).square().sum(-1),
|
||||
"joint_torques_l2": leg_torque.square().sum(-1) + wheel_torque.square().sum(-1),
|
||||
"dof_pos_limits": limits,
|
||||
"terminated": 1. - survival_prob,
|
||||
}
|
||||
|
||||
if constants.get("rolling_rewards", False) and not constants.get("wheel_local_rewards", False):
|
||||
if scan is None:
|
||||
raise ValueError("Rolling rewards require a terrain scan")
|
||||
from .rolling_rewards import rolling_costs
|
||||
terms.update(rolling_costs(state, command, scan, constants["default_q"]))
|
||||
if constants.get("obstacle_progress", False):
|
||||
from .obstacle_rewards import obstacle_progress
|
||||
terms["obstacle_progress"] = obstacle_progress(state, command, scan)
|
||||
if constants.get("wheel_body_intrusion", False):
|
||||
from .wheel_intrusion import wheel_body_intrusion
|
||||
terms["wheel_body_intrusion"] = wheel_body_intrusion(
|
||||
state, constants["default_q"])
|
||||
if constants.get("wheel_local_rewards", False):
|
||||
from .wheel_terrain_rewards import wheel_terms
|
||||
terms.update(wheel_terms(state, command, scan, constants["default_q"]))
|
||||
if constants.get("bad_orientation_penalty", False):
|
||||
from .upward import bad_orientation_value
|
||||
terms["bad_orientation_penalty"] = bad_orientation_value(gravity)
|
||||
from .upward import apply_reward_clamps
|
||||
apply_reward_clamps(terms, constants)
|
||||
return terms
|
||||
|
||||
|
||||
def lambda_returns(rewards, continuation, next_values, gamma=0.99, lam=0.95):
|
||||
carry = next_values[-1]
|
||||
result = []
|
||||
for t in reversed(range(len(rewards))):
|
||||
carry = rewards[t] + gamma * continuation[t] * ((1 - lam) * next_values[t] + lam * carry)
|
||||
result.append(carry)
|
||||
return torch.stack(result[::-1])
|
||||
|
||||
|
||||
def discounted_objective(rewards, continuation, terminal_value, gamma=.99):
|
||||
alive = torch.ones_like(terminal_value)
|
||||
total = torch.zeros_like(terminal_value)
|
||||
discount = 1.
|
||||
for reward, cont in zip(rewards, continuation):
|
||||
total = total + discount * alive * reward
|
||||
alive = alive * cont
|
||||
discount *= gamma
|
||||
return total + discount * alive * terminal_value
|
||||
|
||||
|
||||
def make_imagined_step(model, actor, constants, weights, env_action_order,
|
||||
dt=.02, noise=True, field=None):
|
||||
"""Pure single-step graph, shared by eager and compiled execution."""
|
||||
def step(state, previous, history, pose, scan, command, context, bias, timers):
|
||||
action = actor(history)
|
||||
model_context = context if field is None else torch.cat((context, scan), -1)
|
||||
collision_prob = None
|
||||
if constants.get('nonwheel_collision', False):
|
||||
nxt, delta, _, collision_logit = model.forward_all(state, action, previous, model_context)
|
||||
collision_prob = collision_logit.squeeze(-1).sigmoid()
|
||||
else:
|
||||
nxt, delta = model(state, action, previous, model_context)
|
||||
if field is not None:
|
||||
from .terrain import advance_pose
|
||||
pose = advance_pose(pose, delta)
|
||||
next_scan = None if field is None else field.scan(nxt, pose)
|
||||
cont = continuation_value(nxt, constants)
|
||||
if collision_prob is not None:
|
||||
cont = cont * (1. - collision_prob)
|
||||
terms = reward_terms(nxt, action, previous, command, constants, cont,
|
||||
scan=next_scan)
|
||||
if constants.get('pair_body_intrusion', False):
|
||||
from .wheel_intrusion import pair_body_intrusion
|
||||
terms['pair_body_intrusion'], timers = pair_body_intrusion(
|
||||
nxt, constants['default_q'], timers, dt)
|
||||
reward = dt * sum(weights[name] * value for name, value in terms.items())
|
||||
latest = frame(nxt, action[:, env_action_order], command, bias, noise=noise,
|
||||
wheel_vel=actor.actor_wheel_vel)
|
||||
history = push_history(history, latest, actor.history, actor.terms)
|
||||
return nxt, action, history, pose, next_scan, reward, cont, timers
|
||||
return step
|
||||
|
||||
|
||||
def rollout(model, actor, target_critic, seed, constants, weights, env_action_order,
|
||||
horizon=16, dt=0.02, gamma=0.99, lam=0.95, noise=True, field=None,
|
||||
step_fn=None):
|
||||
"""Actor: full discounted return + terminal value. Critic: detached TD(lambda).
|
||||
|
||||
Caller freezes dynamics and target critic weights, preserving state gradients.
|
||||
Reuse one terrain scan at each state, including the rollout's initial state.
|
||||
"""
|
||||
state, previous, history = seed["state"], seed["previous"], seed["history"]
|
||||
command, context, bias = seed["command"], seed["context"], seed["bias"]
|
||||
pose = seed.get("world_pose")
|
||||
timers = seed.get("pair_timer")
|
||||
if timers is None:
|
||||
timers = torch.zeros(state.shape[0], 4, device=state.device)
|
||||
def value_input(s, prev, scan):
|
||||
base = critic_input(s, command, context, prev)
|
||||
return base if field is None else torch.cat((base, scan), -1)
|
||||
if step_fn is None:
|
||||
step_fn = make_imagined_step(model, actor, constants, weights, env_action_order,
|
||||
dt, noise, field)
|
||||
rewards, continuations, inputs, next_inputs = [], [], [], []
|
||||
scan = None if field is None else field.scan(state, pose)
|
||||
for _ in range(horizon):
|
||||
inputs.append(value_input(state, previous, scan).detach())
|
||||
state, previous, history, pose, scan, reward, cont, timers = step_fn(
|
||||
state, previous, history, pose, scan, command, context, bias, timers)
|
||||
rewards.append(reward)
|
||||
continuations.append(cont)
|
||||
next_inputs.append(value_input(state, previous, scan).detach())
|
||||
final_value = target_critic(value_input(state, previous, scan)).squeeze(-1)
|
||||
r, c = torch.stack(rewards), torch.stack(continuations)
|
||||
objective = discounted_objective(r, c, final_value, gamma)
|
||||
with torch.no_grad():
|
||||
next_values = target_critic(torch.stack(next_inputs)).squeeze(-1)
|
||||
targets = lambda_returns(r.detach(), c.detach(), next_values, gamma, lam)
|
||||
return -objective.mean(), torch.stack(inputs), targets
|
||||
@@ -0,0 +1,52 @@
|
||||
"""Explicit straight/side/yaw command modes, with more stair traversal samples.
|
||||
|
||||
``deadzone`` zeroes per-axis commands below the threshold; v7 uses 0.05 so tiny
|
||||
creep commands become standing instead of ambiguous near-zero motion.
|
||||
"""
|
||||
from dataclasses import dataclass
|
||||
import torch
|
||||
from mjlab.tasks.velocity.mdp.velocity_command import UniformVelocityCommand, UniformVelocityCommandCfg
|
||||
|
||||
|
||||
def mix_commands(values, stairs):
|
||||
result = values.clone()
|
||||
mode = torch.rand(len(values),device=values.device)
|
||||
# General terrain: 40% pure vx, 20% pure vy, 20% pure yaw, 20% combined.
|
||||
result[mode<.4,1:] = 0.
|
||||
result[(mode>=.4)&(mode<.6),0] = 0.
|
||||
result[(mode>=.4)&(mode<.6),2] = 0.
|
||||
result[(mode>=.6)&(mode<.8),:2] = 0.
|
||||
# Stairs: 80% low-speed pure forward/backward, 20% full-range commands.
|
||||
result[stairs] = values[stairs]
|
||||
traverse = stairs & (mode<.8)
|
||||
magnitude = .2 + .4*torch.rand(len(values),device=values.device)
|
||||
result[traverse,0] = torch.where(values[traverse,0]<0,-magnitude[traverse],magnitude[traverse])
|
||||
result[traverse,1:] = 0.
|
||||
return result
|
||||
|
||||
|
||||
class MixedVelocityCommand(UniformVelocityCommand):
|
||||
def _resample_command(self, env_ids):
|
||||
super()._resample_command(env_ids)
|
||||
gen = self._env.cfg.scene.terrain.terrain_generator
|
||||
stair_columns = torch.tensor([c.kind in ('up','down') for c in gen.sub_terrains.values()],device=self.device)
|
||||
stairs = stair_columns[self._env.scene.terrain.terrain_types[env_ids]]
|
||||
mixed = mix_commands(self.vel_command_b[env_ids],stairs)
|
||||
# Stair magnitudes must obey the current curriculum bounds, including reverse.
|
||||
for i, bounds in enumerate((self.cfg.ranges.lin_vel_x, self.cfg.ranges.lin_vel_y, self.cfg.ranges.ang_vel_z)):
|
||||
mixed[:,i].clamp_(*bounds)
|
||||
# Per-axis deadband: tiny commands are standing, not an ambiguous creep.
|
||||
if self.cfg.deadzone > 0.:
|
||||
mixed = torch.where(mixed.abs() < self.cfg.deadzone,
|
||||
torch.zeros_like(mixed), mixed)
|
||||
self.vel_command_b[env_ids] = mixed
|
||||
if self.cfg.hold_stair_command:
|
||||
self.time_left[env_ids[stairs]] = self._env.cfg.episode_length_s + 1.
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class MixedVelocityCommandCfg(UniformVelocityCommandCfg):
|
||||
hold_stair_command: bool = False
|
||||
deadzone: float = 0.
|
||||
def build(self, env):
|
||||
return MixedVelocityCommand(self,env)
|
||||
@@ -0,0 +1,34 @@
|
||||
"""Bounded directional motion near uneven ground, shared by real and imagined steps."""
|
||||
import torch
|
||||
from .world_model import LIN_VEL
|
||||
|
||||
|
||||
def obstacle_gate(scan):
|
||||
local = scan.reshape(-1, 21, 13)[:, 5:16, 3:10].flatten(1)
|
||||
return ((local.amax(-1) - local.amin(-1) - .02) / .08).clamp(0., 1.)
|
||||
|
||||
|
||||
def obstacle_progress(state, command, scan):
|
||||
speed = command[:, :2].norm(dim=-1)
|
||||
along = (state[:, LIN_VEL][:, :2] * command[:, :2]).sum(-1) / speed.clamp_min(.1)
|
||||
# Signed and capped: no benefit for wheel spin, pure yaw, lifting or overspeed.
|
||||
return obstacle_gate(scan) * (speed > .1) * (along / speed.clamp_min(.2)).clamp(-1., 1.)
|
||||
|
||||
|
||||
def env_obstacle_progress(env):
|
||||
from .terrain import TerrainField, TerrainBridge
|
||||
from .stair_curriculum import buffers
|
||||
bridge = getattr(env, '_s10_posture_bridge', None)
|
||||
if bridge is None:
|
||||
bridge = TerrainBridge(env, TerrainField.from_env(env))
|
||||
env._s10_posture_bridge = bridge
|
||||
state = bridge.state()
|
||||
scan = bridge.field.scan(state, bridge.pose())
|
||||
command = env.command_manager.get_command('twist')
|
||||
buffers(env)
|
||||
stalled = ((command[:, :2].norm(dim=-1) > .1)
|
||||
& (state[:, LIN_VEL][:, :2].norm(dim=-1) < .05)
|
||||
& (obstacle_gate(scan) > .5))
|
||||
env._stair_stall_time = torch.where(stalled, env._stair_stall_time + env.step_dt, 0.)
|
||||
env._stair_stuck |= env._stair_stall_time >= 3.
|
||||
return obstacle_progress(state, command, scan)
|
||||
@@ -0,0 +1,149 @@
|
||||
"""S10 actor history and actor/critic networks, independent of the quadruped learner.
|
||||
|
||||
Two Actor observation schemas coexist:
|
||||
|
||||
* v4 / v4-upward / new v7: 53 values/frame, including measured wheel speed;
|
||||
* legacy v7: 49 values/frame, no measured wheel speed (the last action still
|
||||
contains the four wheel commands).
|
||||
|
||||
The learner, replay and deployment exporter select the schema explicitly. Old
|
||||
checkpoints are recognized from their first-layer width.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from .world_model import RunningNorm, ANG_VEL, GRAVITY, LEG_POS, LEG_VEL, WHEEL_VEL
|
||||
|
||||
LEGACY_TERMS = (3, 3, 12, 12, 4, 16, 3)
|
||||
TERMS = (3, 3, 12, 12, 16, 3)
|
||||
LEGACY_FRAME_DIM = sum(LEGACY_TERMS)
|
||||
FRAME_DIM = sum(TERMS)
|
||||
CRITIC_DIM = 38 + 3 + 7 + 16
|
||||
POLICY_FORMAT = "s10-flat-policy-v1"
|
||||
|
||||
|
||||
def frame(state, last_env_action, command, bias, noise=False, wheel_vel=True):
|
||||
"""Exactly the term order/scaling in env.flat (noise before scaling)."""
|
||||
values = [state[:, ANG_VEL], state[:, GRAVITY], state[:, LEG_POS] + bias,
|
||||
state[:, LEG_VEL]]
|
||||
widths = [0.2, 0.05, 0.01, 1.5]
|
||||
if wheel_vel:
|
||||
values.append(state[:, WHEEL_VEL])
|
||||
widths.append(1.5)
|
||||
values.extend((last_env_action, command))
|
||||
widths.extend((0., 0.))
|
||||
if noise:
|
||||
values = [v + (torch.rand_like(v) * 2 - 1) * w if w else v
|
||||
for v, w in zip(values, widths)]
|
||||
# Leg velocity is always index 3; wheel velocity, when present, is index 4.
|
||||
velocity_indices = (3, 4) if wheel_vel else (3,)
|
||||
return torch.cat([v * (0.05 if i in velocity_indices else 1.)
|
||||
for i, v in enumerate(values)], -1)
|
||||
|
||||
|
||||
def push_history(history, latest, length, terms):
|
||||
"""Term-wise histories, oldest to newest, matching mjlab flattening."""
|
||||
result, hidx, fidx = [], 0, 0
|
||||
for dim in terms:
|
||||
result.append(torch.cat((history[:, hidx + dim:hidx + dim * length],
|
||||
latest[:, fidx:fidx + dim]), -1))
|
||||
hidx += dim * length
|
||||
fidx += dim
|
||||
return torch.cat(result, -1)
|
||||
|
||||
|
||||
def critic_input(state, command, context, previous):
|
||||
return torch.cat((state, command, context, previous), -1)
|
||||
|
||||
|
||||
class Network(nn.Module):
|
||||
def __init__(self, inputs, outputs, hidden):
|
||||
super().__init__()
|
||||
self.norm = RunningNorm(inputs)
|
||||
layers = []
|
||||
for n in hidden:
|
||||
layers.extend((nn.Linear(inputs, n), nn.ELU()))
|
||||
inputs = n
|
||||
layers.append(nn.Linear(inputs, outputs))
|
||||
self.net = nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
return self.net(self.norm(x))
|
||||
|
||||
|
||||
class Actor(Network):
|
||||
def __init__(self, history=5, hidden=(256, 128, 64), actor_wheel_vel=True):
|
||||
self.actor_wheel_vel = bool(actor_wheel_vel)
|
||||
self.terms = LEGACY_TERMS if self.actor_wheel_vel else TERMS
|
||||
self.frame_dim = sum(self.terms)
|
||||
super().__init__(self.frame_dim * history, 16, hidden)
|
||||
self.history, self.hidden = history, tuple(hidden)
|
||||
nn.init.uniform_(self.net[-1].weight, -0.01, 0.01)
|
||||
nn.init.zeros_(self.net[-1].bias)
|
||||
|
||||
def forward(self, obs):
|
||||
# Same bounded action distribution during collection and imagination.
|
||||
return super().forward(obs).tanh()
|
||||
|
||||
|
||||
def checkpoint_wheel_vel(checkpoint):
|
||||
"""Infer legacy observation schema from weights; validate explicit metadata."""
|
||||
flag = checkpoint.get("actor_wheel_vel")
|
||||
if "actor" in checkpoint:
|
||||
dim = checkpoint["actor"]["net.0.weight"].shape[1]
|
||||
history = checkpoint["history"]
|
||||
if dim not in (FRAME_DIM * history, LEGACY_FRAME_DIM * history):
|
||||
raise ValueError("Unsupported actor observation dimension")
|
||||
inferred = dim == LEGACY_FRAME_DIM * history
|
||||
if flag is not None and bool(flag) != inferred:
|
||||
raise ValueError("Actor observation metadata disagrees with weights")
|
||||
return inferred
|
||||
return bool(flag) if flag is not None else False
|
||||
|
||||
|
||||
def load_actor(checkpoint, device="cpu"):
|
||||
if checkpoint.get("format") not in (POLICY_FORMAT, "s10-terrain-policy-v1"):
|
||||
raise ValueError("Expected a S10 policy checkpoint, not other-robot or dynamics-only weights")
|
||||
actor = Actor(checkpoint["history"], checkpoint["actor_hidden"],
|
||||
checkpoint_wheel_vel(checkpoint)).to(device)
|
||||
actor.load_state_dict(checkpoint["actor"])
|
||||
return actor.eval()
|
||||
|
||||
|
||||
def evaluation_env_cfg(checkpoint, level=0, kind=None):
|
||||
"""Select the matching task; terrain remains privileged at deployment."""
|
||||
from s10_supermimic.tasks import get_task
|
||||
actor_wheel_vel = checkpoint_wheel_vel(checkpoint)
|
||||
if checkpoint.get("format") == "s10-terrain-policy-v1":
|
||||
checkpoint_task = checkpoint.get("config", {}).get("task")
|
||||
task = get_task(checkpoint_task) if checkpoint_task else None
|
||||
if task is not None:
|
||||
if task.terrain_revision != checkpoint["config"].get("terrain_revision", 1):
|
||||
raise ValueError("Checkpoint task revision differs from its config")
|
||||
if task.actor.wheel_vel != actor_wheel_vel:
|
||||
raise ValueError("Checkpoint task Actor schema differs from its weights")
|
||||
cfg = task.build_env_cfg(play=True, history=checkpoint["history"],
|
||||
seed=checkpoint["config"]["seed"])
|
||||
else:
|
||||
from s10_supermimic.env.terrain import s10_velocity_terrain_env_cfg
|
||||
cfg = s10_velocity_terrain_env_cfg(
|
||||
play=True, history=checkpoint["history"],
|
||||
seed=checkpoint["config"]["seed"],
|
||||
revision=checkpoint["config"].get("terrain_revision", 1),
|
||||
actor_wheel_vel=actor_wheel_vel)
|
||||
# A single staircase can overlap the robot's pre-reset origin at compile
|
||||
# time. Reserve contacts for that transient pose as well as normal rolling.
|
||||
cfg.sim.nconmax = max(cfg.sim.nconmax or 0, 512)
|
||||
cfg.sim.njmax = max(cfg.sim.njmax or 0, 2048)
|
||||
if level not in range(6):
|
||||
raise ValueError("Terrain level must be 0..5")
|
||||
gen = cfg.scene.terrain.terrain_generator
|
||||
# Evaluation fixes difficulty; one row, same seeded generator.
|
||||
gen.num_rows = 1
|
||||
gen.difficulty_range = (level/5,level/5)
|
||||
if kind is not None:
|
||||
gen.sub_terrains = {kind: gen.sub_terrains[kind]}
|
||||
return cfg
|
||||
from s10_supermimic.env.flat import s10_velocity_flat_env_cfg
|
||||
return s10_velocity_flat_env_cfg(play=True, history=checkpoint["history"],
|
||||
actor_wheel_vel=actor_wheel_vel)
|
||||
@@ -0,0 +1,21 @@
|
||||
"""Isolate invalid simulation rows before terrain lookup or replay insertion."""
|
||||
import torch
|
||||
|
||||
|
||||
def invalid_rows(env, obs):
|
||||
d = env.scene["robot"].data
|
||||
tensors = (d.root_link_pos_w, d.projected_gravity_b, d.root_link_lin_vel_b,
|
||||
d.root_link_ang_vel_b, d.joint_pos, d.joint_vel, obs["actor"])
|
||||
return ~torch.stack([torch.isfinite(t).flatten(1).all(-1) for t in tensors]).all(0)
|
||||
|
||||
|
||||
def recover_invalid_envs(env, obs, reward):
|
||||
bad = invalid_rows(env, obs) | ~torch.isfinite(reward)
|
||||
ids = bad.nonzero(as_tuple=False).flatten()
|
||||
if ids.numel():
|
||||
print(f"RECOVERY reset {ids.numel()} invalid environments; first IDs={ids[:16].tolist()}", flush=True)
|
||||
obs, _ = env.reset(env_ids=ids)
|
||||
# A failed reset must be handled by the supervisor, never added to replay.
|
||||
if invalid_rows(env, obs).any():
|
||||
raise RuntimeError("Environment reset failed to restore finite state")
|
||||
return obs, bad
|
||||
@@ -0,0 +1,157 @@
|
||||
"""Bounded S10 transition storage with chronological multi-step sampling."""
|
||||
|
||||
import torch
|
||||
from .world_model import STATE_DIM, ACTION_DIM, CONTEXT_DIM
|
||||
|
||||
|
||||
class Replay:
|
||||
def __init__(self, length, num_envs, device, history_dim=0, terrain=False,
|
||||
contacts=False, courses=False, collisions=False, pairs=False):
|
||||
if min(length, num_envs) < 1:
|
||||
raise ValueError("Replay length and environment count must be positive")
|
||||
self.length, self.size, self.cursor = length, 0, 0
|
||||
self.archive = None
|
||||
self.success_archive = None
|
||||
dims = dict(state=STATE_DIM, action=ACTION_DIM, previous=ACTION_DIM,
|
||||
context=CONTEXT_DIM, next=STATE_DIM, pose=3, valid=1)
|
||||
if terrain:
|
||||
dims["world_pose"] = 3
|
||||
if contacts:
|
||||
dims["next_contacts"] = 16
|
||||
if collisions:
|
||||
dims.update(collision=1, collision_valid=1)
|
||||
if courses:
|
||||
dims["course"] = 4 # type, level, uphill ground transition, top crossing
|
||||
if pairs:
|
||||
dims["pair_timer"] = 4 # per-wheel continuous 5/6-region intrusion time
|
||||
if history_dim:
|
||||
dims.update(history=history_dim, command=3, bias=12)
|
||||
self.data = {k: torch.zeros(length, num_envs, dim, device=device) for k, dim in dims.items()}
|
||||
|
||||
def add(self, **items):
|
||||
if items.keys() != self.data.keys():
|
||||
raise ValueError("Replay transition fields do not match configured fields")
|
||||
for k, value in items.items():
|
||||
self.data[k][self.cursor].copy_(value.detach())
|
||||
self.cursor = (self.cursor + 1) % self.length
|
||||
self.size = min(self.size + 1, self.length)
|
||||
|
||||
def sample(self, batch, window, keys=None):
|
||||
if self.size < window or window < 1:
|
||||
raise ValueError("Not enough transitions for a window")
|
||||
device = self.data["state"].device
|
||||
start = torch.randint(self.size - window + 1, (batch,), device=device)
|
||||
oldest = self.cursor if self.size == self.length else 0
|
||||
times = (oldest + start[:, None] + torch.arange(window, device=device)) % self.length
|
||||
envs = torch.randint(self.data["state"].shape[1], (batch, 1), device=device)
|
||||
return {k: v[times, envs] for k, v in self.data.items() if keys is None or k in keys}
|
||||
|
||||
|
||||
def sample_world_model(self, batch, window, prioritized=.25, archived=.25,
|
||||
capacity=100000, keep_prob=.02):
|
||||
keys = ('state','next','action','previous','context','pose','valid',
|
||||
'world_pose','next_contacts')
|
||||
if 'next_contacts' not in self.data:
|
||||
return self.sample(batch, window, keys)
|
||||
if self.archive is None:
|
||||
self.archive = WindowArchive(capacity)
|
||||
if 'course' in self.data:
|
||||
return self._sample_courses(batch, window, capacity, keep_prob)
|
||||
n_priority = int(batch * prioritized)
|
||||
n_archive = int(batch * archived) if self.archive.size else 0
|
||||
parts = [self.sample(batch-n_priority-n_archive, window, keys)]
|
||||
candidates = self.sample(4*max(n_priority,1), window, keys)
|
||||
score, rare = event_scores(candidates)
|
||||
selected = rare & (torch.rand_like(score) < keep_prob)
|
||||
self.archive.add({k:v[selected] for k,v in candidates.items()})
|
||||
if n_priority:
|
||||
indices = (score + .001*torch.rand_like(score)).topk(n_priority).indices
|
||||
parts.append({k:v[indices] for k,v in candidates.items()})
|
||||
if n_archive:
|
||||
parts.append(self.archive.sample(n_archive))
|
||||
return {k:torch.cat([p[k] for p in parts]) for k in parts[0]}
|
||||
|
||||
|
||||
def _sample_courses(self, batch, window, capacity, keep_prob):
|
||||
# Total archive capacity remains bounded by --archive-size, split evenly.
|
||||
if self.success_archive is None:
|
||||
self.archive = WindowArchive(max(1, capacity // 2))
|
||||
self.success_archive = WindowArchive(max(1, capacity - capacity // 2))
|
||||
keys = tuple(k for k in self.data if k not in ('history', 'command', 'bias'))
|
||||
n_priority = batch // 4
|
||||
candidates = self.sample(max(1, n_priority * 4), window, keys)
|
||||
score, rare = event_scores(candidates)
|
||||
valid = candidates['valid'].squeeze(-1).bool().cumprod(1).bool()
|
||||
tag = candidates['course']
|
||||
moving = (candidates['pose'][:, :, :2].norm(dim=-1) > .001) & valid
|
||||
up = tag[:, :, 0] == 1
|
||||
useful = ((tag[:, :, 2:] > .5).any(-1) & valid & up).any(-1)
|
||||
# Do not fill the general archive with repeated stationary stair impacts.
|
||||
rare &= (~up.any(-1) | moving.any(-1))
|
||||
if 'collision' in candidates:
|
||||
from .collision import collision_events
|
||||
rare |= collision_events(candidates)
|
||||
for archive, chosen in ((self.archive, rare), (self.success_archive, useful)):
|
||||
chosen = chosen & (torch.rand_like(score) < keep_prob)
|
||||
archive.add({k:v[chosen] for k,v in candidates.items()})
|
||||
n_success = batch // 8 if self.success_archive.size else 0
|
||||
n_rare = batch // 8 if self.archive.size else 0
|
||||
# At least half the batch is always uniform replay, including when sparse.
|
||||
parts = [self.sample(batch - n_priority - n_success - n_rare, window, keys)]
|
||||
if n_priority:
|
||||
rank = score + 8 * useful.float() + 2 * (moving & up).any(-1).float()
|
||||
selected = (rank + .001 * torch.rand_like(rank)).topk(n_priority).indices
|
||||
parts.append({k:v[selected] for k,v in candidates.items()})
|
||||
if n_success: parts.append(self.success_archive.sample(n_success))
|
||||
if n_rare: parts.append(self.archive.sample(n_rare))
|
||||
return {k:torch.cat([part[k] for part in parts]) for k in keys}
|
||||
|
||||
|
||||
def event_scores(w):
|
||||
"""Measured wheel contact transitions; ignore all labels after a reset.
|
||||
|
||||
Force impacts, sustained flight, hard vertical motion and episode ends enter
|
||||
the archive. Contact changes also include normal wheel touchdown/liftoff.
|
||||
"""
|
||||
valid = w['valid'].squeeze(-1).bool()
|
||||
mask = valid.cumprod(1).bool()
|
||||
contacts = w['next_contacts'][:,:,:4] > .5
|
||||
pairs = mask[:,1:] & mask[:,:-1]
|
||||
changes = ((contacts[:,1:] != contacts[:,:-1]).sum(-1)*pairs).sum(-1)
|
||||
flight = ((~contacts.any(-1)) & mask).sum(-1)
|
||||
force = torch.nan_to_num(w['next_contacts'][:,:,4:],nan=0.).reshape(*mask.shape,4,3).norm(dim=-1)
|
||||
impact = ((force.amax(-1)>200.) & mask).any(-1)
|
||||
hard_vz = ((w['state'][:,:,6].abs()>.6) & mask).any(-1)
|
||||
# A terminal endpoint is not a valid dynamics target, but its preceding
|
||||
# window is still useful. Windows starting after termination are excluded.
|
||||
ended = ((~valid) & torch.cat((torch.ones_like(mask[:,:1]),mask[:,:-1]),1)).any(-1)
|
||||
score = changes.float() + flight.float() + 4.*ended + 2.*impact
|
||||
rare = ((flight>=2) | hard_vz | impact | ended) & mask[:,0]
|
||||
score = score * mask[:,0]
|
||||
if 'collision' in w:
|
||||
from .collision import collision_events
|
||||
collision = collision_events(w)
|
||||
score = score + 8.*collision
|
||||
rare |= collision
|
||||
return score, rare
|
||||
|
||||
|
||||
class WindowArchive:
|
||||
"""Bounded GPU ring storing only world-model fields, not actor histories."""
|
||||
def __init__(self, capacity):
|
||||
self.capacity, self.size, self.cursor = capacity, 0, 0
|
||||
self.data = None
|
||||
|
||||
def add(self, windows):
|
||||
n = min(len(windows['state']),self.capacity)
|
||||
if not n:return
|
||||
if self.data is None:
|
||||
self.data = {k:v.new_empty((self.capacity,*v.shape[1:])) for k,v in windows.items()}
|
||||
ids = (torch.arange(n,device=windows['state'].device)+self.cursor)%self.capacity
|
||||
for k,v in windows.items():self.data[k][ids] = v[-n:].detach()
|
||||
self.cursor = (self.cursor+n)%self.capacity
|
||||
self.size = min(self.capacity,self.size+n)
|
||||
|
||||
def sample(self, batch):
|
||||
ids = torch.randint(self.size,(batch,),device=self.data['state'].device)
|
||||
return {k:v[ids] for k,v in self.data.items()}
|
||||
@@ -0,0 +1,77 @@
|
||||
"""Terrain-gated rolling preference, shared by real and imagined rewards."""
|
||||
import torch
|
||||
from .world_model import GRAVITY, LEG_POS, LEG_VEL
|
||||
|
||||
|
||||
def rolling_gate(scan, command):
|
||||
local = scan.reshape(-1, 21, 13)[:, 5:16, 3:10].flatten(1)
|
||||
relief = local.amax(-1) - local.amin(-1)
|
||||
# Flat (<2 cm): favor rolling; near obstacles (>=6 cm): allow lifting.
|
||||
flat = (1. - (relief - .02) / .04).clamp(0., 1.)
|
||||
# Non-steerable wheels may require stepping for lateral/yaw maneuvers.
|
||||
straight = torch.exp(-(command[:,1]/.25).square() - (command[:,2]/1.).square())
|
||||
return flat * straight
|
||||
|
||||
|
||||
def wheel_geometry(state, default_q):
|
||||
"""S10 wheel centres in body coordinates and vertical cylinder half-extents.
|
||||
|
||||
Joint order is hipx (abduction), hipy (flexion), knee. Geometry follows
|
||||
``S10_description/s10_mjcf/S10.xml``: hipx at x=±0.2277/y=±0.0574, hipy
|
||||
offset y=±0.0779, knee at z=-0.18, wheel at y=±0.045891/z=-0.18,
|
||||
wheel radius 0.081 m and half-width 0.0185 m. The analytic result was
|
||||
checked against MuJoCo ``mj_forward`` to machine precision.
|
||||
|
||||
The cylinder is symmetric about its axle, so wheel phase is unnecessary.
|
||||
"""
|
||||
q = (state[:,LEG_POS] + default_q).reshape(-1,4,3)
|
||||
hipx, hipy, knee = q.unbind(-1)
|
||||
side = state.new_tensor([1.,-1.,1.,-1.])
|
||||
# Leg sagittal plane: knee link (0.18 m) + wheel link (0.18 m).
|
||||
x = .18*knee.sin()
|
||||
z = -.18 - .18*knee.cos()
|
||||
x3 = hipy.cos()*x - hipy.sin()*z
|
||||
z3 = hipy.sin()*x + hipy.cos()*z
|
||||
y3 = side*(.0779+.045891)
|
||||
y4 = hipx.cos()*y3 + hipx.sin()*z3
|
||||
z4 = -hipx.sin()*y3 + hipx.cos()*z3
|
||||
point = torch.stack((state.new_tensor([.2277,.2277,-.2277,-.2277]) + x3,
|
||||
state.new_tensor([.0574,-.0574,.0574,-.0574]) + y4,
|
||||
z4),-1)
|
||||
up = -torch.nn.functional.normalize(state[:,GRAVITY],dim=-1,eps=1e-6)
|
||||
# Wheel axle is local -y, rotated only by hipx; sign is irrelevant after abs.
|
||||
axle_z = -up[:,1:2]*hipx.cos()+up[:,2:3]*hipx.sin()
|
||||
radius_z = .081*(1.-axle_z.square()).clamp_min(1e-8).sqrt()+.0185*axle_z.abs()
|
||||
return point, radius_z
|
||||
|
||||
|
||||
def wheel_bottom_height(state, default_q):
|
||||
"""Wheel bottoms relative to the ground below the base.
|
||||
|
||||
Only used on near-flat ground by rolling_gate, not as obstacle clearance.
|
||||
"""
|
||||
point, radius_z = wheel_geometry(state, default_q)
|
||||
up = -torch.nn.functional.normalize(state[:,GRAVITY],dim=-1,eps=1e-6)
|
||||
return state[:,:1] + (point*up[:,None,:]).sum(-1) - radius_z
|
||||
|
||||
|
||||
def rolling_costs(state, command, scan, default_q):
|
||||
gate = rolling_gate(scan, command)
|
||||
# Tolerate 5 mm of clearance; no positive reward for lifting anywhere.
|
||||
lift = ((wheel_bottom_height(state,default_q)-.005).clamp_min(0)/.05).square().sum(-1)
|
||||
# Reference mjlab task uses mean squared speed and relaxes during recovery.
|
||||
tilt = (state[:,GRAVITY][:,:2].norm(dim=-1)-.08).div(.22).clamp(0.,1.)
|
||||
motion = state[:,LEG_VEL].square().mean(-1) * (1.-.8*tilt)
|
||||
return {'unnecessary_wheel_lift':gate*lift, 'rolling_leg_motion':gate*motion}
|
||||
|
||||
|
||||
def env_rolling_cost(env, term):
|
||||
from .terrain import TerrainField, TerrainBridge
|
||||
bridge = getattr(env,'_s10_posture_bridge',None)
|
||||
if bridge is None:
|
||||
bridge = TerrainBridge(env,TerrainField.from_env(env))
|
||||
env._s10_posture_bridge = bridge
|
||||
state = bridge.state()
|
||||
return rolling_costs(state, env.command_manager.get_command('twist'),
|
||||
bridge.field.scan(state,bridge.pose()),
|
||||
bridge.robot.data.default_joint_pos[0,bridge.legs])[term]
|
||||
@@ -0,0 +1,75 @@
|
||||
"""V5 episode outcomes and performance-gated velocity curriculum.
|
||||
|
||||
All buffers are training-only metadata, never actor observations.
|
||||
"""
|
||||
import torch
|
||||
|
||||
KINDS = ('flat', 'up', 'down', 'platform', 'slope', 'rubble')
|
||||
MIN_ITERATIONS = (0, 3000, 8000, 15000, 22000)
|
||||
|
||||
|
||||
def buffers(env):
|
||||
if not hasattr(env, '_stair_stats'):
|
||||
env._stair_stats = torch.zeros(6, 6, 4, device=env.device) # trials, pass, fall, stuck
|
||||
env._stair_stall_time = torch.zeros(env.num_envs, device=env.device)
|
||||
env._stair_stuck = torch.zeros(env.num_envs, dtype=torch.bool, device=env.device)
|
||||
return env._stair_stats
|
||||
|
||||
|
||||
def stair_pass(offset, gravity, level, up):
|
||||
"""Require body/rear clearance of the last riser, correct landing height and stability."""
|
||||
heights = offset.new_tensor([.02, .04, .08, .12, .16, .20])[level]
|
||||
expected = torch.where(up, 5 * heights, -5 * heights)
|
||||
return ((offset[:, 0].abs() > 2.8) & (offset[:, 1].abs() < 2.5)
|
||||
& (offset[:, 2] > expected + .20) & (offset[:, 2] < expected + .60)
|
||||
& (gravity[:, :2].norm(dim=-1) < .4))
|
||||
|
||||
|
||||
def curriculum(env, env_ids):
|
||||
stats = buffers(env)
|
||||
terrain = env.scene.terrain
|
||||
d = env.scene['robot'].data
|
||||
offset = d.root_link_pos_w[env_ids] - env.scene.env_origins[env_ids]
|
||||
levels = terrain.terrain_levels[env_ids].clone()
|
||||
types = terrain.terrain_types[env_ids]
|
||||
# Six fixed columns, generated in KINDS order.
|
||||
stairs = (types == 1) | (types == 2)
|
||||
active = env.episode_length_buf[env_ids] > 0
|
||||
fallen = env.termination_manager.terminated[env_ids]
|
||||
success = (offset[:, 0].abs() > 5.4) & (offset[:, 1].abs() < 6.5)
|
||||
top = stair_pass(offset, d.projected_gravity_b[env_ids], levels, types == 1)
|
||||
# Landing contacts rule out simply jumping through the clearance test.
|
||||
contact = env.scene['wheel_ground_contact'].data.found[env_ids].reshape(len(env_ids), -1).gt(0).sum(-1) >= 2
|
||||
success = torch.where(stairs, top & contact, success) & active & ~fallen
|
||||
command = env.command_manager.get_command('twist')[env_ids]
|
||||
traversal = (command[:, 0].abs() > .1) & (command[:, 1:].abs().amax(-1) < .05)
|
||||
trials = (~stairs | traversal) & active & ((env.episode_length_buf[env_ids] * env.step_dt >= 5.) | fallen)
|
||||
values = torch.stack((trials, success & trials, fallen & trials,
|
||||
env._stair_stuck[env_ids] & trials), -1).float()
|
||||
stats.view(-1, 4).index_add_(0, types * 6 + levels, values)
|
||||
terrain.terrain_levels[env_ids] = (levels + success.long() - (active & fallen).long()).clamp(0, 5)
|
||||
terrain.env_origins[env_ids] = terrain.terrain_origins[terrain.terrain_levels[env_ids], types]
|
||||
env._stair_stall_time[env_ids] = 0
|
||||
env._stair_stuck[env_ids] = False
|
||||
return terrain.terrain_levels.float().mean()
|
||||
|
||||
|
||||
def next_stage(stage, iteration, stats, velocity_mae):
|
||||
"""Minimum age AND sufficient actual stair passes, never time alone."""
|
||||
if stage >= 4 or iteration < MIN_ITERATIONS[stage + 1]:
|
||||
return stage
|
||||
# Stage increases require success at 4, 8, 12, then 16 cm respectively.
|
||||
trials, passed = stats[1, stage + 1, :2]
|
||||
accurate = max(velocity_mae[:2]) < .18 and velocity_mae[2] < .30
|
||||
return stage + int(trials >= 64 and passed / max(trials, 1) >= .60 and accurate)
|
||||
|
||||
|
||||
def summary(env):
|
||||
stats = buffers(env).detach().cpu().tolist()
|
||||
levels, types = env.scene.terrain.terrain_levels, env.scene.terrain.terrain_types
|
||||
result = {}
|
||||
for i, name in enumerate(KINDS):
|
||||
selected = levels[types == i]
|
||||
result[name] = dict(level_counts=torch.bincount(selected, minlength=6).cpu().tolist(),
|
||||
episodes_by_level=stats[i])
|
||||
return result
|
||||
@@ -0,0 +1,95 @@
|
||||
"""Shared static height atlas and privileged, differentiable yaw-frame scans."""
|
||||
import numpy as np
|
||||
import torch
|
||||
from .bridge import S10Bridge
|
||||
|
||||
SCAN_DIM = 21 * 13
|
||||
|
||||
class TerrainField:
|
||||
def __init__(self, heights, origin, resolution=.04):
|
||||
self.h, self.origin, self.res = heights, origin, resolution
|
||||
x,y = torch.meshgrid(torch.linspace(-1,1,21,device=heights.device),
|
||||
torch.linspace(-.6,.6,13,device=heights.device), indexing="ij")
|
||||
self.offsets = torch.stack((x.flatten(),y.flatten()),-1)
|
||||
|
||||
def sample(self, xy):
|
||||
# Clamp *integer* neighbour indices too. NaN survives float clamp and
|
||||
# casts to INT64_MIN; on CUDA that would poison the entire device context.
|
||||
finite = torch.isfinite(xy).all(-1)
|
||||
clean = torch.where(torch.isfinite(xy), xy, torch.zeros_like(xy))
|
||||
gx = ((clean[...,0]-self.origin[0])/self.res).clamp(0,self.h.shape[1]-1)
|
||||
gy = ((clean[...,1]-self.origin[1])/self.res).clamp(0,self.h.shape[0]-1)
|
||||
x = gx.floor().long().clamp(0,self.h.shape[1]-2)
|
||||
y = gy.floor().long().clamp(0,self.h.shape[0]-2)
|
||||
a,b = gx-x,gy-y
|
||||
value = ((1-a)*self.h[y,x]+a*self.h[y,x+1])*(1-b)+((1-a)*self.h[y+1,x]+a*self.h[y+1,x+1])*b
|
||||
# Do not disguise corrupt states as valid ground: the learner will discard
|
||||
# this update, while real invalid environments are reset before querying.
|
||||
return torch.where(finite, value, torch.full_like(value, float("nan")))
|
||||
|
||||
def scan(self, state, pose):
|
||||
c,s = pose[:,2:3].cos(),pose[:,2:3].sin()
|
||||
dx,dy = self.offsets[:,0],self.offsets[:,1]
|
||||
xy = pose[:,None,:2]+torch.stack((c*dx-s*dy,s*dx+c*dy),-1)
|
||||
return (state[:,:1]+self.sample(pose[:,:2])[:,None]-self.sample(xy)).clamp(-1.5,1.5)
|
||||
|
||||
@classmethod
|
||||
def from_env(cls, env, resolution=.04):
|
||||
"""Rasterize exact oriented boxes in small bounding rectangles, no per-pixel rays.
|
||||
|
||||
Shared CPU build is ~6 MB at 4 cm for a 48 x 48 m atlas. No robot geometry.
|
||||
Only these static box courses are supported; fail if new shapes are introduced.
|
||||
"""
|
||||
import mujoco
|
||||
model = env.sim.mj_model
|
||||
data = mujoco.MjData(model)
|
||||
mujoco.mj_forward(model,data)
|
||||
robot_ids = set(env.scene["robot"].indexing.geom_ids.cpu().tolist())
|
||||
ids = [i for i in range(model.ngeom) if i not in robot_ids and model.geom_contype[i]]
|
||||
corners = np.array([[a,b,c] for a in (-1,1) for b in (-1,1) for c in (-1,1)])
|
||||
boxes=[]
|
||||
for i in ids:
|
||||
if model.geom_type[i] != mujoco.mjtGeom.mjGEOM_BOX:
|
||||
raise ValueError(f"Unsupported terrain geometry: {model.geom(i).name}")
|
||||
center, rot, size = data.geom_xpos[i],data.geom_xmat[i].reshape(3,3),model.geom_size[i]
|
||||
points = (corners*size) @ rot.T+center
|
||||
boxes.append((center.copy(),rot.copy(),size.copy(),points.min(0),points.max(0)))
|
||||
lo=np.min([b[3][:2] for b in boxes],axis=0)-resolution
|
||||
hi=np.max([b[4][:2] for b in boxes],axis=0)+resolution
|
||||
shape=np.ceil((hi-lo)/resolution).astype(int)+1
|
||||
heights=np.zeros((shape[1],shape[0]),dtype=np.float32)
|
||||
for center,rot,size,low,high in boxes:
|
||||
start=np.maximum(0,np.floor((low[:2]-lo)/resolution).astype(int))
|
||||
end=np.minimum(shape,np.ceil((high[:2]-lo)/resolution).astype(int)+1)
|
||||
xs=lo[0]+np.arange(start[0],end[0])*resolution
|
||||
ys=lo[1]+np.arange(start[1],end[1])*resolution
|
||||
x,y=np.meshgrid(xs,ys)
|
||||
# Intersect vertical line p=(x,y,z) with all three oriented-box slabs.
|
||||
lower=np.full_like(x,-np.inf); upper=np.full_like(x,np.inf); valid=np.ones_like(x,dtype=bool)
|
||||
for k in range(3):
|
||||
base=rot[0,k]*(x-center[0])+rot[1,k]*(y-center[1])-rot[2,k]*center[2]
|
||||
slope=rot[2,k]
|
||||
if abs(slope)<1e-10:
|
||||
valid &= np.abs(base)<=size[k]+1e-9
|
||||
else:
|
||||
z1=(-size[k]-base)/slope; z2=(size[k]-base)/slope
|
||||
lower=np.maximum(lower,np.minimum(z1,z2)); upper=np.minimum(upper,np.maximum(z1,z2))
|
||||
surface=np.where(valid & (lower<=upper),upper,0)
|
||||
patch=heights[start[1]:end[1],start[0]:end[0]]
|
||||
np.maximum(patch,surface,out=patch)
|
||||
return cls(torch.as_tensor(heights,device=env.device),tuple(lo),resolution)
|
||||
|
||||
|
||||
def advance_pose(pose, delta):
|
||||
c,s=pose[:,2].cos(),pose[:,2].sin()
|
||||
return pose+torch.stack((c*delta[:,0]-s*delta[:,1],s*delta[:,0]+c*delta[:,1],delta[:,2]),-1)
|
||||
|
||||
class TerrainBridge(S10Bridge):
|
||||
def __init__(self, env, field):
|
||||
super().__init__(env)
|
||||
self.field=field
|
||||
|
||||
def state(self):
|
||||
state=super().state()
|
||||
state[:,0]-=self.field.sample(self.robot.data.root_link_pos_w[:,:2])
|
||||
return state
|
||||
@@ -0,0 +1,33 @@
|
||||
"""Shared real/imagined posture objective for the opt-in terrain v2 task.
|
||||
|
||||
S10 joint order is hipx (abduction), hipy (flexion), knee. Hipx keeps a 2x
|
||||
weight while hipy/knee relax on rough terrain. v7 sets
|
||||
``relief_floor=0.02`` so stairs leave knee/hipy deviation almost unpenalised
|
||||
and the policy can flex instead of holding a stiff leg.
|
||||
"""
|
||||
import torch
|
||||
|
||||
LEG_SCALES = (0.25, 0.6, 0.8) * 4
|
||||
|
||||
|
||||
def posture_cost(q, scan, relief_floor=0.125):
|
||||
# Only the supporting/approaching ground near the body should relax posture.
|
||||
# scan is body height minus ground; its range is independent of body height.
|
||||
local = scan.reshape(-1, 21, 13)[:, 5:16, 3:10].flatten(1)
|
||||
relief = ((local.amax(-1) - local.amin(-1)) / 0.12).clamp(0., 1.)
|
||||
multiplier = (1. - relief).clamp_min(relief_floor)
|
||||
legs = q.reshape(-1, 4, 3).square()
|
||||
# Hip abduction stays constrained even on steps; flexion is freer there.
|
||||
return 2. * legs[:, :, 0].sum(-1) + multiplier * legs[:, :, 1:].sum((1, 2))
|
||||
|
||||
|
||||
def env_posture_cost(env, asset_cfg, relief_floor=0.125):
|
||||
from .terrain import TerrainField, TerrainBridge
|
||||
bridge = getattr(env, "_s10_posture_bridge", None)
|
||||
if bridge is None:
|
||||
bridge = TerrainBridge(env, TerrainField.from_env(env))
|
||||
env._s10_posture_bridge = bridge
|
||||
d = env.scene[asset_cfg.name].data
|
||||
q = d.joint_pos[:, asset_cfg.joint_ids] - d.default_joint_pos[:, asset_cfg.joint_ids]
|
||||
return posture_cost(q, bridge.field.scan(bridge.state(), bridge.pose()),
|
||||
relief_floor)
|
||||
@@ -0,0 +1,120 @@
|
||||
"""S10 terrain dynamics. Flat model checkpoints remain a separate format."""
|
||||
import torch
|
||||
from torch import nn
|
||||
from .world_model import S10WorldModel, RunningNorm
|
||||
from .terrain import SCAN_DIM, advance_pose
|
||||
|
||||
FORMAT = "s10-terrain-dynamics-v1"
|
||||
POLICY_FORMAT = "s10-terrain-policy-v1"
|
||||
|
||||
class TerrainWorldModel(S10WorldModel):
|
||||
def __init__(self, hidden=(512,512,512), contacts=False, collisions=False):
|
||||
super().__init__(hidden)
|
||||
self.net[0] = nn.Linear(77+SCAN_DIM,hidden[0])
|
||||
self.input_norm = RunningNorm(77+SCAN_DIM)
|
||||
self.predict_contacts = contacts
|
||||
if contacts:
|
||||
self.contact_head = nn.Linear(hidden[-1],16)
|
||||
self.predict_collisions = collisions
|
||||
if collisions:
|
||||
self.collision_head = nn.Linear(hidden[-1],1)
|
||||
nn.init.zeros_(self.collision_head.weight)
|
||||
nn.init.constant_(self.collision_head.bias, -4.6)
|
||||
self.last_contact_metrics = {}
|
||||
|
||||
def forward_all(self, state, action, previous_action, context):
|
||||
"""Predict endpoint state, pose, optional contacts and non-wheel collision logit."""
|
||||
x = self.input_norm(self.inputs(state,action,previous_action,context))
|
||||
for layer in list(self.net.children())[:-1]:
|
||||
x = layer(x)
|
||||
delta = self.net[-1](x)*self.output_norm.scale+self.output_norm.mean
|
||||
contacts = self.contact_head(x) if self.predict_contacts else None
|
||||
collision = self.collision_head(x) if self.predict_collisions else None
|
||||
return state+delta[:,:38],delta[:,38:],contacts,collision
|
||||
|
||||
def forward_with_contacts(self, state, action, previous_action, context):
|
||||
"""Predict endpoint state, pose, and optional [logits4, asinh(force/100)12]."""
|
||||
pred,delta,contacts,_ = self.forward_all(state,action,previous_action,context)
|
||||
return pred,delta,contacts
|
||||
|
||||
def forward(self, state, action, previous_action, context):
|
||||
pred,delta,_ = self.forward_with_contacts(state,action,previous_action,context)
|
||||
return pred,delta
|
||||
|
||||
@staticmethod
|
||||
def decode_contacts(prediction):
|
||||
"""Contact probabilities and signed base-frame force vectors in newtons."""
|
||||
return prediction[:,:4].sigmoid(),100.*prediction[:,4:].clamp(-10,10).sinh().reshape(-1,4,3)
|
||||
|
||||
@staticmethod
|
||||
def inputs(state, action, previous_action, context):
|
||||
if context.shape[-1] != 7+SCAN_DIM:
|
||||
raise ValueError("Terrain dynamics require context7 plus height scan273")
|
||||
return torch.cat((S10WorldModel.inputs(state,action,previous_action,context[:,:7]),context[:,7:]),-1)
|
||||
|
||||
def terrain_loss(self, b, field):
|
||||
state,pose=b["state"][:,0],b["world_pose"][:,0]
|
||||
mask=b["valid"].squeeze(-1).bool().cumprod(1).bool()
|
||||
# Include the first terminal transition for collision classification only.
|
||||
prefix=torch.cat((torch.ones_like(mask[:,:1]),mask[:,:-1]),1)
|
||||
total=next(self.parameters()).sum()*0.
|
||||
contact_total=total.detach().clone(); force_total=total.detach().clone()
|
||||
collision_total=total.clone(); collision_count=total.detach().clone()
|
||||
positives=total.detach().clone(); predicted=total.detach().clone()
|
||||
for t in range(b["action"].shape[1]):
|
||||
keep=prefix[:,t] if self.predict_collisions else mask[:,t]
|
||||
if not keep.any(): break
|
||||
context=torch.cat((b["context"][keep,0],field.scan(state[keep],pose[keep])),-1)
|
||||
pred,delta,contact,collision=self.forward_all(state[keep],b["action"][keep,t],b["previous"][keep,t],context)
|
||||
current=mask[keep,t]
|
||||
# Never regress the reset state as the endpoint of a collision.
|
||||
if current.any():
|
||||
error=torch.cat((pred[current]-b["next"][keep,t][current],delta[current]-b["pose"][keep,t][current]),-1)
|
||||
total=total+(error/self.output_norm.scale).square().mean(-1).sum()
|
||||
if self.predict_contacts:
|
||||
target=b["next_contacts"][keep,t][current]
|
||||
finite=torch.isfinite(target).all(-1)
|
||||
# Missing/non-finite sensor labels must not poison dynamics gradients.
|
||||
prediction=contact[current][finite]; target=target[finite]
|
||||
classification=torch.nn.functional.binary_cross_entropy_with_logits(
|
||||
prediction[:,:4],target[:,:4],reduction='none').mean(-1)
|
||||
force=torch.nn.functional.smooth_l1_loss(prediction[:,4:],
|
||||
torch.asinh(target[:,4:]/100.),reduction='none').mean(-1)
|
||||
total=total+.1*(classification+force).sum()
|
||||
contact_total=contact_total+classification.detach().sum()
|
||||
force_total=force_total+force.detach().sum()
|
||||
if self.predict_collisions:
|
||||
label=b['collision'][keep,t]
|
||||
usable=b['collision_valid'][keep,t].squeeze(-1).bool() & torch.isfinite(label).all(-1)
|
||||
logits=collision[usable]; labels=label[usable]
|
||||
loss=torch.nn.functional.binary_cross_entropy_with_logits(logits,labels,reduction='sum')
|
||||
collision_total=collision_total+loss
|
||||
collision_count=collision_count+usable.sum()
|
||||
positives=positives+labels.detach().sum()
|
||||
predicted=predicted+logits.detach().sigmoid().sum()
|
||||
next_pose=advance_pose(pose[keep],delta)
|
||||
state=state.clone(); pose=pose.clone()
|
||||
state[keep]=pred; pose[keep]=next_pose
|
||||
count=mask.sum().clamp_min(1)
|
||||
if self.predict_contacts:
|
||||
self.last_contact_metrics={'contact_bce':contact_total/count,'contact_force_loss':force_total/count}
|
||||
result=total/count
|
||||
if self.predict_collisions:
|
||||
denom=collision_count.clamp_min(1)
|
||||
result=result+collision_total/denom
|
||||
self.last_contact_metrics.update(collision_bce=collision_total.detach()/denom,
|
||||
collision_label_rate=positives/denom,collision_prediction_mean=predicted/denom,
|
||||
collision_label_count=collision_count)
|
||||
return result
|
||||
|
||||
def checkpoint(self):
|
||||
return {"format":FORMAT,"hidden":self.hidden,"contacts":self.predict_contacts,
|
||||
"collisions":self.predict_collisions,"model":self.state_dict()}
|
||||
|
||||
@classmethod
|
||||
def from_checkpoint(cls, checkpoint, device="cpu"):
|
||||
if checkpoint.get("format") != FORMAT: raise ValueError("Expected terrain dynamics checkpoint")
|
||||
model=cls(checkpoint["hidden"],contacts=checkpoint.get("contacts",False),
|
||||
collisions=checkpoint.get("collisions",False)).to(device)
|
||||
model.load_state_dict(checkpoint["model"])
|
||||
return model
|
||||
@@ -0,0 +1,782 @@
|
||||
"""Online S10 model-based learning: collect, fit dynamics, imagine, optimize."""
|
||||
|
||||
import argparse
|
||||
import copy
|
||||
import json
|
||||
import math
|
||||
import time
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from zoneinfo import ZoneInfo
|
||||
|
||||
import torch
|
||||
|
||||
from .bridge import S10Bridge, pose_delta
|
||||
from .world_model import S10WorldModel
|
||||
from .policy import (Actor, Network, CRITIC_DIM, POLICY_FORMAT,
|
||||
checkpoint_wheel_vel, critic_input)
|
||||
from .replay import Replay
|
||||
from .imagine import frozen, rollout, make_imagined_step
|
||||
from .recovery import recover_invalid_envs
|
||||
from .upward import (TILT_STAGES_DEG, init_upward_state,
|
||||
upward_curriculum_stats, upward_video_tilt_limit)
|
||||
from s10_supermimic.tasks import get_task, list_tasks, task_names
|
||||
|
||||
|
||||
def _format_value(value):
|
||||
if value is None:
|
||||
return "n/a"
|
||||
if isinstance(value, float):
|
||||
return f"{value:+.4f}" if abs(value) >= 1e-3 else f"{value:.3e}"
|
||||
if isinstance(value, (list, tuple)):
|
||||
return "[" + ", ".join(_format_value(v) for v in value) + "]"
|
||||
return str(value)
|
||||
|
||||
|
||||
def format_training_log(row, args, reward_terms, termination_rates):
|
||||
"""IsaacLab-style multi-line training summary; the JSONL stays machine-readable."""
|
||||
def line(label, value):
|
||||
return f"{label:>42s}: {value}"
|
||||
total = row.get("total_transitions", 0)
|
||||
elapsed = max(row.get("elapsed_s", 0.0), 1e-9)
|
||||
lines = [
|
||||
"",
|
||||
"=" * 78,
|
||||
f"Learning iteration {row.get('step', 0)}/{args.iterations}",
|
||||
line("Total steps", f"{total:,}"),
|
||||
line("Steps per second", f"{total / elapsed:,.0f}"),
|
||||
line("Iteration time", f"{row.get('iteration_time_s', 0.0):.3f}s"),
|
||||
line("Collection time", f"{row.get('collection_time_s', 0.0):.3f}s"),
|
||||
line("Learning time", f"{row.get('learning_time_s', 0.0):.3f}s"),
|
||||
line("Mean reward", f"{row.get('reward', 0.0):+.4f}"),
|
||||
line("Mean episode length", f"{row.get('mean_episode_length', 0.0):.1f}"),
|
||||
line("Actor loss", _format_value(row.get("actor_loss"))),
|
||||
line("Critic loss", _format_value(row.get("critic_loss"))),
|
||||
line("World model loss", _format_value(row.get("wm_loss"))),
|
||||
line("Mean action noise", f"{args.action_noise:.3f}"),
|
||||
line("Skipped updates", row.get("skipped_updates", 0)),
|
||||
"-" * 78,
|
||||
]
|
||||
for name, value in reward_terms.items():
|
||||
lines.append(line(f"Episode_Reward/{name}", f"{value:+.4f}"))
|
||||
for key in ("terrain_level_mean", "terrain_level_frac",
|
||||
"tilt_level_mean", "tilt_level_q75", "tilt_level_frac",
|
||||
"tilt_limit_mean_deg", "recovery_window_rate",
|
||||
"traverse_window_rate", "recovery_window_ready"):
|
||||
if key in row:
|
||||
lines.append(line(f"Curriculum/{key}", _format_value(row[key])))
|
||||
for key in ("command_stage", "command_mode", "command_ranges"):
|
||||
if key in row:
|
||||
lines.append(line(f"Command/{key}", _format_value(row[key])))
|
||||
vmae = row.get("velocity_mae")
|
||||
if vmae:
|
||||
lines.append(line("Metrics/base_velocity/error_vel_xy",
|
||||
f"{(vmae[0] + vmae[1]) / 2:.4f}"))
|
||||
lines.append(line("Metrics/base_velocity/error_vel_yaw", f"{vmae[2]:.4f}"))
|
||||
for name, value in termination_rates.items():
|
||||
lines.append(line(f"Episode_Termination/{name}", f"{value:.4f}"))
|
||||
lines.append("=" * 78)
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def constants_from_bridge(bridge):
|
||||
"""Fixed gains/limits of this S10 task; no per-environment gain randomization."""
|
||||
d, dev = bridge.robot.data, bridge.env.device
|
||||
from .terrain_posture import env_posture_cost, LEG_SCALES
|
||||
from .upward import UPWARD_REWARD_CLAMPS, is_upward_task
|
||||
from s10_supermimic.robot.config import LEG_JOINT_NAMES
|
||||
upward = is_upward_task(bridge.env)
|
||||
revised = (bridge.env.cfg.rewards["leg_posture"].func is env_posture_cost
|
||||
or "wheel_obstacle_clearance" in bridge.env.cfg.rewards
|
||||
or upward)
|
||||
# Read the actual leg actuator gains so the imagined torque surrogate matches
|
||||
# v4 (kp=80) and the softer v7 (kp=20) without duplicating constants.
|
||||
robot_cfg = bridge.env.cfg.scene.entities["robot"]
|
||||
leg_act = next(
|
||||
(a for a in robot_cfg.articulation.actuators
|
||||
if set(LEG_JOINT_NAMES).issubset(set(getattr(a, "target_names_expr", ())))),
|
||||
None)
|
||||
leg_kp = float(getattr(leg_act, "stiffness", 80.0))
|
||||
leg_kd = float(getattr(leg_act, "damping", 2.0))
|
||||
leg_effort = float(getattr(leg_act, "effort_limit", 50.0))
|
||||
return {
|
||||
"terrain_posture_v2": revised,
|
||||
"nonwheel_collision": "illegal_nonwheel_contact" in bridge.env.cfg.terminations,
|
||||
"wheel_local_rewards": "wheel_obstacle_clearance" in bridge.env.cfg.rewards,
|
||||
"rolling_rewards": "unnecessary_wheel_lift" in bridge.env.cfg.rewards,
|
||||
"obstacle_progress": "obstacle_progress" in bridge.env.cfg.rewards,
|
||||
"wheel_body_intrusion": "wheel_body_intrusion" in bridge.env.cfg.rewards,
|
||||
"pair_body_intrusion": "pair_body_intrusion" in bridge.env.cfg.rewards,
|
||||
"posture_relief_floor": bridge.env.cfg.rewards["leg_posture"].params.get(
|
||||
"relief_floor", 0.125),
|
||||
"upright_mode": "upward" if upward else "exp",
|
||||
"bad_orientation_penalty": upward,
|
||||
"reward_clamps": dict(UPWARD_REWARD_CLAMPS) if upward else {},
|
||||
"wheel_action_scale": bridge.env.cfg.actions["wheel_vel"].scale,
|
||||
"leg_action_scale": torch.tensor(LEG_SCALES if revised else (.25,) * 12, device=dev),
|
||||
"terrain_rewards": hasattr(bridge, "field"),
|
||||
"upright_std": bridge.env.cfg.rewards["upright"].params.get("std", .6),
|
||||
"default_q": d.default_joint_pos[0, bridge.legs].detach().clone(),
|
||||
"limits": d.soft_joint_pos_limits[0, bridge.legs].detach().clone(),
|
||||
# S10 leg actuators: v4 kp=80/kd=2, v7 kp=20/kd=2; effort=50 N·m.
|
||||
# TODO(s10): the -2e-5 joint_torques_l2 weight was inherited from the
|
||||
# initial recipe; retune if S10 training shows torque suppression or chatter.
|
||||
"kp": torch.full((12,), leg_kp, device=dev),
|
||||
"kd": torch.full((12,), leg_kd, device=dev),
|
||||
"effort": torch.full((12,), leg_effort, device=dev),
|
||||
# S10 wheel actuator: velocity damping 0.8, effort 14 N·m, 35 rad/s clip.
|
||||
"wheel_kd": 0.8,
|
||||
"wheel_effort": 14.0,
|
||||
"wheel_velocity_limit": 35.0,
|
||||
}
|
||||
|
||||
|
||||
def actor_learning_rate(args, iteration):
|
||||
age = max(0, iteration - args.actor_schedule_start)
|
||||
progress = min(1., max(0, age - args.actor_hold_iterations) / max(1, args.lr_decay_iterations))
|
||||
decay = args.lr_final_frac + (1.-args.lr_final_frac)*.5*(1.+math.cos(math.pi*progress))
|
||||
ramp = min(1., (age+1) / max(1, args.actor_ramp_iterations))
|
||||
return args.lr * args.actor_lr_scale * decay * ramp
|
||||
|
||||
|
||||
def update(optimizer, loss, parameters, element_clip=None):
|
||||
parameters = list(parameters)
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
if not torch.isfinite(loss):
|
||||
print("RECOVERY skipped non-finite loss; optimizer unchanged", flush=True)
|
||||
return False
|
||||
loss.backward()
|
||||
norm = torch.nn.utils.clip_grad_norm_(parameters, 1., error_if_nonfinite=False)
|
||||
if not torch.isfinite(norm):
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
print("RECOVERY skipped non-finite gradient; optimizer unchanged", flush=True)
|
||||
return False
|
||||
if element_clip is not None:
|
||||
torch.nn.utils.clip_grad_value_(parameters, element_clip)
|
||||
optimizer.step()
|
||||
return True
|
||||
|
||||
|
||||
def record_policy_video(actor, args, output, iteration, env_order,
|
||||
upward_tilt_limit=None, command_ranges=None):
|
||||
"""Record a short MP4 with the current actor on a fresh single-env scene.
|
||||
|
||||
This mirrors Isaac Lab's ``--video``/``--video_length``/``--video_interval``
|
||||
workflow. A separate one-env scene is used so recording does not perturb the
|
||||
training environments. Recording is skipped if the optional mediapy/EGL
|
||||
dependencies or the extra scene cannot be created.
|
||||
"""
|
||||
import random
|
||||
import numpy as np
|
||||
|
||||
state = (
|
||||
random.getstate(),
|
||||
np.random.get_state(),
|
||||
torch.get_rng_state(),
|
||||
torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None,
|
||||
)
|
||||
video_env = None
|
||||
try:
|
||||
import mediapy
|
||||
from mjlab.envs import ManagerBasedRlEnv
|
||||
from .bridge import S10Bridge
|
||||
|
||||
task = get_task(args.task) if getattr(args, "task", None) else None
|
||||
if task is not None:
|
||||
if task.actor.wheel_vel != actor.actor_wheel_vel:
|
||||
raise ValueError("Task Actor schema differs from the training Actor")
|
||||
cfg = task.build_env_cfg(play=True, history=args.history, seed=args.seed)
|
||||
elif args.terrain:
|
||||
from s10_supermimic.env.terrain import s10_velocity_terrain_env_cfg
|
||||
cfg = s10_velocity_terrain_env_cfg(
|
||||
play=True, history=args.history, seed=args.seed,
|
||||
revision=args.terrain_revision, actor_wheel_vel=actor.actor_wheel_vel)
|
||||
else:
|
||||
from s10_supermimic.env.flat import s10_velocity_flat_env_cfg
|
||||
cfg = s10_velocity_flat_env_cfg(
|
||||
play=True, history=args.history,
|
||||
actor_wheel_vel=actor.actor_wheel_vel)
|
||||
cfg.scene.num_envs = 1
|
||||
cfg.rewards = {}
|
||||
cfg.metrics = {}
|
||||
cfg.observations.pop("critic", None)
|
||||
cfg.episode_length_s = 1e9
|
||||
cfg.viewer.width = args.video_width
|
||||
cfg.viewer.height = args.video_height
|
||||
cfg.viewer.distance = 2.5
|
||||
cfg.viewer.elevation = -20.
|
||||
cfg.viewer.azimuth = 140.
|
||||
video_env = ManagerBasedRlEnv(
|
||||
cfg=cfg, device=args.device, render_mode="rgb_array")
|
||||
if command_ranges is not None:
|
||||
ranges = video_env.command_manager.get_term("twist").cfg.ranges
|
||||
ranges.lin_vel_x, ranges.lin_vel_y, ranges.ang_vel_z = command_ranges
|
||||
bridge = S10Bridge(video_env)
|
||||
if list(bridge.to_env) != list(env_order):
|
||||
raise ValueError("Video env action order differs from the training env")
|
||||
if upward_tilt_limit is not None:
|
||||
video_env._upward_tilt_limits = torch.full(
|
||||
(video_env.num_envs,), float(upward_tilt_limit), device=video_env.device)
|
||||
obs, _ = video_env.reset()
|
||||
path = output / "videos" / f"iteration_{iteration:06d}.mp4"
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
was_training = actor.training
|
||||
actor.eval()
|
||||
with torch.no_grad(), mediapy.VideoWriter(
|
||||
str(path), shape=(args.video_height, args.video_width),
|
||||
fps=args.video_fps) as writer:
|
||||
for _ in range(args.video_length):
|
||||
action = actor(obs["actor"])
|
||||
obs, _, terminated, truncated, _ = video_env.step(
|
||||
bridge.env_action(action))
|
||||
writer.add_image(video_env.render())
|
||||
if bool(terminated.any() or truncated.any()):
|
||||
obs, _ = video_env.reset()
|
||||
actor.train(was_training)
|
||||
print(f"S10 video saved: {path}", flush=True)
|
||||
finally:
|
||||
if video_env is not None:
|
||||
video_env.close()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
random.setstate(state[0])
|
||||
np.random.set_state(state[1])
|
||||
torch.set_rng_state(state[2])
|
||||
if state[3] is not None:
|
||||
torch.cuda.set_rng_state_all(state[3])
|
||||
|
||||
|
||||
class Learner:
|
||||
"""Pure-Torch learner; also usable in small CPU tests without mjlab."""
|
||||
def __init__(self, args, field=None):
|
||||
self.args, self.field = args, field
|
||||
self.format = POLICY_FORMAT
|
||||
self.model_class = S10WorldModel
|
||||
extra = 0
|
||||
if field is not None:
|
||||
from .terrain_world_model import TerrainWorldModel, POLICY_FORMAT as TERRAIN_FORMAT
|
||||
from .terrain import SCAN_DIM
|
||||
self.model_class, self.format, extra = TerrainWorldModel, TERRAIN_FORMAT, SCAN_DIM
|
||||
self.model = self.model_class(tuple(args.wm_hidden), **(
|
||||
{"contacts":args.terrain_revision >= 4,
|
||||
"collisions":args.terrain_revision >= 7} if field is not None else {})).to(args.device)
|
||||
self.actor = Actor(args.history, tuple(args.actor_hidden),
|
||||
args.actor_wheel_vel).to(args.device)
|
||||
self.critic = Network(CRITIC_DIM + extra, 1, tuple(args.critic_hidden)).to(args.device)
|
||||
self.target = copy.deepcopy(self.critic)
|
||||
self.optimizers = [torch.optim.Adam(module.parameters(), lr=args.lr)
|
||||
for module in (self.model, self.actor, self.critic)]
|
||||
self.policy_updates = 0
|
||||
self.skipped_updates = 0
|
||||
self.terrain_state = None
|
||||
self.curriculum_state = None
|
||||
self.env_steps = 0
|
||||
self.imagined_step = None
|
||||
|
||||
def critic_input(self, state, command, context, previous, pose):
|
||||
x = critic_input(state, command, context, previous)
|
||||
return x if self.field is None else torch.cat((x, self.field.scan(state, pose)), -1)
|
||||
|
||||
def fit(self, b):
|
||||
if self.field is not None:
|
||||
loss = self.model.terrain_loss(b, self.field)
|
||||
if not update(self.optimizers[0], loss, self.model.parameters()):
|
||||
self.skipped_updates += 1
|
||||
return None
|
||||
return loss.item()
|
||||
loss = self.model.loss(b["state"][:, 0], b["action"], b["previous"], b["context"][:, 0],
|
||||
b["next"], b["pose"], b["valid"].squeeze(-1))
|
||||
if not update(self.optimizers[0], loss, self.model.parameters()):
|
||||
self.skipped_updates += 1
|
||||
return None
|
||||
return loss.item()
|
||||
|
||||
def improve(self, seed, constants, weights, env_order, dt):
|
||||
args = self.args
|
||||
if args.compile_imagination and self.imagined_step is None:
|
||||
step = make_imagined_step(self.model,self.actor,constants,weights,env_order,
|
||||
dt=dt,field=self.field)
|
||||
self.imagined_step = torch.compile(step, dynamic=False, fullgraph=True)
|
||||
with frozen(self.model, self.target):
|
||||
actor_loss, inputs, targets = rollout(
|
||||
self.model, self.actor, self.target, seed, constants, weights, env_order,
|
||||
horizon=args.horizon, dt=dt, gamma=args.gamma, lam=args.lam, field=self.field,
|
||||
step_fn=self.imagined_step)
|
||||
if not update(self.optimizers[1], actor_loss, self.actor.parameters(), args.actor_grad_clip_element):
|
||||
self.skipped_updates += 1
|
||||
return None, None
|
||||
x, y = inputs.flatten(0, 1), targets.flatten()
|
||||
critic_losses = []
|
||||
for _ in range(args.critic_epochs):
|
||||
order = torch.randperm(len(x),device=x.device)
|
||||
for indices in order.tensor_split(min(args.critic_minibatches,len(x))):
|
||||
value = self.critic(x[indices]).squeeze(-1)
|
||||
critic_loss = torch.nn.functional.mse_loss(value, y[indices])
|
||||
if not update(self.optimizers[2], critic_loss, self.critic.parameters()):
|
||||
self.skipped_updates += 1
|
||||
return actor_loss.item(), None
|
||||
critic_losses.append(critic_loss.detach())
|
||||
with torch.no_grad():
|
||||
for target, source in zip(self.target.parameters(), self.critic.parameters()):
|
||||
target.lerp_(source, 0.02)
|
||||
# Running normalization is a coordinate system, not an EMA weight.
|
||||
self.target.norm.load_state_dict(self.critic.norm.state_dict())
|
||||
self.policy_updates += 1
|
||||
return actor_loss.item(), torch.stack(critic_losses).mean().item()
|
||||
|
||||
def save(self, path, step, env_order, dt):
|
||||
checkpoint = {
|
||||
"training_algorithm": "s10-full-return-events-v1",
|
||||
"format": self.format, "step": step, "history": self.actor.history,
|
||||
"actor_wheel_vel": self.actor.actor_wheel_vel,
|
||||
"actor_hidden": self.actor.hidden, "actor": self.actor.state_dict(),
|
||||
"world_model": self.model.checkpoint(), "critic": self.critic.state_dict(),
|
||||
"target_critic": self.target.state_dict(), "config": vars(self.args),
|
||||
"optimizers": [o.state_dict() for o in self.optimizers],
|
||||
"terrain_state": self.terrain_state,
|
||||
"curriculum_state": self.curriculum_state,
|
||||
"env_steps": self.env_steps,
|
||||
"policy_updates": self.policy_updates, "env_action_order": list(env_order), "dt": dt,
|
||||
}
|
||||
path = Path(path)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
tmp = path.with_suffix(".tmp")
|
||||
torch.save(checkpoint, tmp)
|
||||
tmp.replace(path)
|
||||
|
||||
def load(self, path, env_order, dt):
|
||||
ckpt = torch.load(path, map_location=self.args.device, weights_only=True)
|
||||
if ckpt.get("format") != self.format:
|
||||
raise ValueError("Resume requires a S10 policy checkpoint")
|
||||
if ckpt["env_action_order"] != list(env_order) or abs(ckpt["dt"] - dt) > 1e-8:
|
||||
raise ValueError("Checkpoint action order/control timestep differs from this environment")
|
||||
if ckpt["history"] != self.args.history:
|
||||
raise ValueError("Checkpoint observation history differs")
|
||||
if self.field is not None and ckpt["config"]["seed"] != self.args.seed:
|
||||
raise ValueError("Terrain resume requires the same terrain seed")
|
||||
if ckpt["config"].get("terrain_revision", 1) != getattr(self.args, "terrain_revision", 1):
|
||||
raise ValueError("Terrain revision/action scale differs; start a new run for a different revision")
|
||||
if checkpoint_wheel_vel(ckpt) != self.actor.actor_wheel_vel:
|
||||
raise ValueError("Actor observation schema differs; choose the matching task or start a new run")
|
||||
self.actor.load_state_dict(ckpt["actor"])
|
||||
self.model.load_state_dict(self.model_class.from_checkpoint(ckpt["world_model"]).state_dict())
|
||||
self.critic.load_state_dict(ckpt["critic"])
|
||||
self.target.load_state_dict(ckpt["target_critic"])
|
||||
for optimizer, data in zip(self.optimizers, ckpt["optimizers"]):
|
||||
optimizer.load_state_dict(data)
|
||||
self.policy_updates = ckpt["policy_updates"]
|
||||
self.terrain_state = ckpt.get("terrain_state")
|
||||
self.curriculum_state = ckpt.get("curriculum_state")
|
||||
self.env_steps = ckpt.get("env_steps", ckpt["step"] * ckpt["config"].get("num_steps_per_env", 1))
|
||||
return ckpt["step"]
|
||||
|
||||
|
||||
def parser():
|
||||
p = argparse.ArgumentParser(description=__doc__)
|
||||
p.add_argument("--task", choices=task_names(), default=None,
|
||||
help="Select a registered S10 task; its defaults are applied before CLI overrides")
|
||||
p.add_argument("--list-tasks", action="store_true", help="List registered tasks and exit")
|
||||
for name, default in (("num-envs", 64), ("iterations", 10000), ("buffer-len", 512),
|
||||
("history", 5), ("wm-window", 8), ("wm-batch", 128),
|
||||
("wm-updates", 2), ("imag-batch", 128), ("horizon", 16),
|
||||
("warmup", 300), ("critic-epochs", 5), ("save-interval", 1000),
|
||||
("log-interval", 50), ("seed", 0), ("num-steps-per-env", 1), ("critic-minibatches", 4),
|
||||
("archive-size", 100000), ("actor-ramp-iterations", 1000),
|
||||
("lr-decay-iterations", 30000), ("actor-hold-iterations", 0),
|
||||
("actor-schedule-start", 0)):
|
||||
p.add_argument("--" + name, type=int, default=default)
|
||||
for name, default in (("lr", 2e-4), ("action-noise", 0.15), ("gamma", 0.99), ("lam", 0.95)):
|
||||
p.add_argument("--" + name, type=float, default=default)
|
||||
p.add_argument("--compile-imagination", action="store_true")
|
||||
p.add_argument("--no-compile-imagination", dest="compile_imagination",
|
||||
action="store_false", default=False)
|
||||
p.add_argument("--prioritized-replay", action="store_true")
|
||||
p.add_argument("--actor-lr-scale", type=float, default=.3)
|
||||
p.add_argument("--lr-final-frac", type=float, default=.03)
|
||||
p.add_argument("--actor-grad-clip-element", type=float, default=.1)
|
||||
p.add_argument("--wm-hidden", type=int, nargs="+", default=[512, 512, 512])
|
||||
p.add_argument("--actor-hidden", type=int, nargs="+", default=[256, 128, 64])
|
||||
p.add_argument("--critic-hidden", type=int, nargs="+", default=[256, 128, 64])
|
||||
p.add_argument("--device", default="cuda:0")
|
||||
p.add_argument("--terrain", action="store_true", help="Blind static rough-terrain task")
|
||||
p.add_argument("--terrain-revision", type=int, choices=(1, 2, 3, 4, 5, 6, 7), default=1,
|
||||
help="2: larger leg actions; 3: bidirectional commands; 4: rolling rewards; 5: performance-gated stair curriculum; 6: wheel-local clearance; 7: safe v7 (53-value Actor, time-ramped 1.5/0.8/1.2 commands, ±45° init, wheel/body intrusion, collision); new run required")
|
||||
p.add_argument("--command-curriculum", action="store_true", help="Progressively widen velocity ranges")
|
||||
p.add_argument("--command-curriculum-start", type=int, default=0, help="Learner iteration where the velocity schedule starts; preserve on resume")
|
||||
p.add_argument("--init-actor", default="", help="Initialize only the actor from a flat policy; dynamics/critic start fresh")
|
||||
p.add_argument("--resume", default="")
|
||||
p.add_argument("--output", default="")
|
||||
p.add_argument("--video", action="store_true", help="Record periodic MP4 videos with the current actor")
|
||||
p.add_argument("--video-length", type=int, default=500, help="Frames per video")
|
||||
p.add_argument("--video-interval", type=int, default=15000, help="Record every N learner iterations")
|
||||
p.add_argument("--video-fps", type=int, default=50, help="Video frame rate")
|
||||
p.add_argument("--video-width", type=int, default=640)
|
||||
p.add_argument("--video-height", type=int, default=480)
|
||||
return p
|
||||
|
||||
|
||||
def main():
|
||||
p = parser()
|
||||
pre, _ = p.parse_known_args()
|
||||
if pre.list_tasks:
|
||||
for spec in list_tasks():
|
||||
print(f"{spec.name}: {spec.description}")
|
||||
return
|
||||
task = get_task(pre.task) if pre.task else None
|
||||
if task is not None:
|
||||
p.set_defaults(**task.defaults)
|
||||
args = p.parse_args()
|
||||
if task is not None and args.terrain_revision != task.terrain_revision:
|
||||
p.error(f"Task {task.name} requires terrain-revision {task.terrain_revision}")
|
||||
if task is not None:
|
||||
args.actor_wheel_vel = task.actor.wheel_vel
|
||||
elif args.resume:
|
||||
resumed = torch.load(args.resume, map_location="cpu", weights_only=True)
|
||||
args.actor_wheel_vel = checkpoint_wheel_vel(resumed)
|
||||
else:
|
||||
args.actor_wheel_vel = args.terrain_revision <= 4 or args.terrain_revision >= 7
|
||||
positive = (args.archive_size, args.critic_minibatches, args.lr_decay_iterations, args.num_steps_per_env, args.num_envs, args.iterations, args.buffer_len, args.history, args.wm_window,
|
||||
args.wm_batch, args.wm_updates, args.imag_batch, args.horizon, args.critic_epochs,
|
||||
args.save_interval, args.log_interval, *args.wm_hidden, *args.actor_hidden,
|
||||
*args.critic_hidden)
|
||||
if min(positive) < 1 or args.warmup < 0 or args.buffer_len < args.wm_window:
|
||||
p.error("Positive sizes required; buffer-len must be >= wm-window; warmup >= 0")
|
||||
if not (0 < args.gamma <= 1 and 0 <= args.lam <= 1 and args.lr > 0 and args.action_noise >= 0):
|
||||
p.error("Invalid optimizer, noise or return parameters")
|
||||
if args.actor_lr_scale <= 0 or not 0 < args.lr_final_frac <= 1 or args.actor_ramp_iterations < 0 or args.actor_grad_clip_element <= 0:
|
||||
p.error("Invalid actor learning-rate/gradient protection")
|
||||
if args.actor_hold_iterations < 0:
|
||||
p.error("actor-hold-iterations must be nonnegative")
|
||||
if args.terrain_revision >= 5 and args.archive_size < 2:
|
||||
p.error("terrain v5+ requires archive-size >=2")
|
||||
if 5 <= args.terrain_revision < 7 and not args.command_curriculum:
|
||||
p.error("terrain v5/v6 requires command-curriculum; v7 uses a time-based cap")
|
||||
if args.video and (args.video_length < 1 or args.video_interval < 1
|
||||
or not 1 <= args.video_fps <= 50
|
||||
or args.video_width < 64 or args.video_height < 64
|
||||
or args.video_width % 2 or args.video_height % 2):
|
||||
p.error("Video requires length/interval >=1, fps 1..50, even width/height >=64")
|
||||
if args.prioritized_replay and (not args.terrain or args.terrain_revision < 4):
|
||||
p.error("prioritized-replay requires measured contacts from terrain revision >=4")
|
||||
from mjlab.envs import ManagerBasedRlEnv
|
||||
from s10_supermimic.env.flat import s10_velocity_flat_env_cfg
|
||||
|
||||
torch.manual_seed(args.seed)
|
||||
if args.resume and args.init_actor:
|
||||
p.error("Choose resume or init-actor")
|
||||
if args.terrain_revision >= 2 and (not args.terrain or args.init_actor):
|
||||
p.error("terrain-revision >=2 requires --terrain and cannot initialize a flat actor")
|
||||
if task is not None:
|
||||
if task.actor.wheel_vel != args.actor_wheel_vel:
|
||||
p.error("Task Actor schema differs from the requested Actor schema")
|
||||
cfg = task.build_env_cfg(play=False, history=args.history, seed=args.seed)
|
||||
elif args.terrain:
|
||||
from s10_supermimic.env.terrain import s10_velocity_terrain_env_cfg
|
||||
cfg = s10_velocity_terrain_env_cfg(
|
||||
history=args.history, seed=args.seed, revision=args.terrain_revision,
|
||||
actor_wheel_vel=args.actor_wheel_vel)
|
||||
else:
|
||||
cfg = s10_velocity_flat_env_cfg(
|
||||
history=args.history, actor_wheel_vel=args.actor_wheel_vel)
|
||||
if args.command_curriculum:
|
||||
if not args.terrain or args.terrain_revision < 4:
|
||||
p.error("command-curriculum requires terrain revision >=4")
|
||||
from .command_curriculum import set_command_stage, stages_for
|
||||
set_command_stage(cfg.commands["twist"], 0,
|
||||
stages=stages_for(args.terrain_revision))
|
||||
cfg.seed = args.seed
|
||||
cfg.scene.num_envs = args.num_envs
|
||||
env = ManagerBasedRlEnv(cfg=cfg, device=args.device)
|
||||
try:
|
||||
if args.task == "s10-v4-upward":
|
||||
init_upward_state(env)
|
||||
obs, _ = env.reset()
|
||||
field = None
|
||||
if args.terrain:
|
||||
from .terrain import TerrainField, TerrainBridge
|
||||
field = TerrainField.from_env(env)
|
||||
bridge = TerrainBridge(env, field)
|
||||
env._s10_posture_bridge = bridge
|
||||
print(f"Shared terrain atlas: {field.h.numel()*4/2**20:.1f} MiB", flush=True)
|
||||
else:
|
||||
bridge = S10Bridge(env)
|
||||
learner = Learner(args, field)
|
||||
if args.init_actor:
|
||||
ckpt = torch.load(args.init_actor, map_location=args.device, weights_only=True)
|
||||
if ckpt.get("format") != POLICY_FORMAT or ckpt["history"] != args.history:
|
||||
raise ValueError("init-actor requires a matching flat S10 actor")
|
||||
if ckpt["env_action_order"] != list(bridge.to_env) or abs(ckpt["dt"]-env.step_dt)>1e-8:
|
||||
raise ValueError("Warm-start action contract differs")
|
||||
if checkpoint_wheel_vel(ckpt) != learner.actor.actor_wheel_vel:
|
||||
raise ValueError("Warm-start actor observation schema differs")
|
||||
learner.actor.load_state_dict(ckpt["actor"])
|
||||
constants = constants_from_bridge(bridge)
|
||||
weights = {name: term.weight for name, term in cfg.rewards.items()}
|
||||
contact_targets = args.terrain and args.terrain_revision >= 4
|
||||
replay = Replay(args.buffer_len, args.num_envs, args.device,
|
||||
learner.actor.frame_dim * args.history,
|
||||
terrain=args.terrain, contacts=contact_targets,
|
||||
courses=args.terrain_revision >= 5,
|
||||
collisions=args.terrain_revision >= 7,
|
||||
pairs=args.terrain_revision >= 7)
|
||||
start = learner.load(args.resume, bridge.to_env, env.step_dt) if args.resume else 0
|
||||
command_stage = 0
|
||||
if args.command_curriculum:
|
||||
saved_stage = None
|
||||
if args.terrain_revision >= 5:
|
||||
saved_stage = (int(learner.curriculum_state['stage'])
|
||||
if learner.curriculum_state else 0)
|
||||
command_stage = set_command_stage(
|
||||
env.command_manager.get_term("twist").cfg,
|
||||
start - args.command_curriculum_start,
|
||||
stage=saved_stage, stages=stages_for(args.terrain_revision))
|
||||
if args.terrain and learner.terrain_state is not None:
|
||||
terrain = env.scene.terrain
|
||||
levels = learner.terrain_state["levels"].to(args.device)
|
||||
types = learner.terrain_state["types"].to(args.device)
|
||||
if levels.shape != terrain.terrain_levels.shape:
|
||||
raise ValueError("Terrain resume requires the same num-envs")
|
||||
terrain.terrain_levels.copy_(levels)
|
||||
terrain.terrain_types.copy_(types)
|
||||
terrain.env_origins.copy_(terrain.terrain_origins[levels, types])
|
||||
obs, _ = env.reset()
|
||||
if args.terrain_revision >= 5:
|
||||
from .stair_curriculum import buffers, next_stage, summary
|
||||
stats = buffers(env)
|
||||
saved = learner.curriculum_state
|
||||
velocity_ema = saved['velocity_ema'] if saved else [1., 1., 1.]
|
||||
if saved:
|
||||
stats.copy_(saved['stats'].to(env.device))
|
||||
if args.terrain_revision >= 5:
|
||||
command_stage = int(saved['stage']) if saved else 0
|
||||
if args.command_curriculum:
|
||||
set_command_stage(env.command_manager.get_term('twist').cfg, 0,
|
||||
stage=command_stage,
|
||||
stages=stages_for(args.terrain_revision))
|
||||
if start >= args.iterations:
|
||||
raise ValueError("iterations must exceed the resumed checkpoint step")
|
||||
if obs["actor"].shape[-1] != learner.actor.frame_dim * args.history:
|
||||
raise ValueError("S10 actor observation dimension mismatch")
|
||||
output = Path(args.output or "logs/s10/" + datetime.now(ZoneInfo("Asia/Shanghai")).strftime("%Y-%m-%d_%H-%M-%S"))
|
||||
output.mkdir(parents=True, exist_ok=True)
|
||||
previous = torch.zeros(args.num_envs, 16, device=args.device)
|
||||
ou = torch.zeros_like(previous)
|
||||
totals = torch.zeros(6, device=args.device)
|
||||
reward_totals = torch.zeros(len(env.reward_manager.active_terms), device=args.device)
|
||||
interval_steps = 0
|
||||
nonwheel_terminations = 0
|
||||
invalid_env_resets = 0
|
||||
episode_length_sum = 0.0
|
||||
episode_count = 0
|
||||
termination_counts = {name: 0 for name in env.termination_manager.active_terms}
|
||||
collection_time_sum = 0.0
|
||||
learning_time_sum = 0.0
|
||||
began = time.monotonic()
|
||||
print(f"S10 replay storage: {sum(t.numel() * t.element_size() for t in replay.data.values()) / 2**20:.1f} MiB", flush=True)
|
||||
if args.video:
|
||||
try:
|
||||
ranges = env.command_manager.get_term("twist").cfg.ranges
|
||||
record_policy_video(
|
||||
learner.actor, args, output, start, bridge.to_env,
|
||||
upward_tilt_limit=(upward_video_tilt_limit(env)
|
||||
if args.task == "s10-v4-upward" else None),
|
||||
command_ranges=(ranges.lin_vel_x, ranges.lin_vel_y, ranges.ang_vel_z))
|
||||
except Exception as exc:
|
||||
print(f"RECOVERY skipped initial video: {exc}", flush=True)
|
||||
with (output / "metrics.jsonl").open("a") as log:
|
||||
for step in range(start, args.iterations):
|
||||
iter_start = time.monotonic()
|
||||
for group in learner.optimizers[1].param_groups:
|
||||
group["lr"] = actor_learning_rate(args, step)
|
||||
if args.command_curriculum and args.terrain_revision < 5:
|
||||
command_stage = set_command_stage(
|
||||
env.command_manager.get_term("twist").cfg,
|
||||
step - args.command_curriculum_start,
|
||||
stages=stages_for(args.terrain_revision))
|
||||
for _ in range(args.num_steps_per_env):
|
||||
with torch.no_grad():
|
||||
history = obs["actor"].clone()
|
||||
state, pose, context = bridge.state(), bridge.pose(), bridge.context()
|
||||
command = env.command_manager.get_command("twist").clone()
|
||||
bias = bridge.robot.data.encoder_bias[:, bridge.legs].clone()
|
||||
learner.actor.norm.update(history)
|
||||
learner.critic.norm.update(learner.critic_input(state, command, context, previous, pose))
|
||||
learner.target.norm.load_state_dict(learner.critic.norm.state_dict())
|
||||
ou = 0.95 * ou + (1 - 0.95**2)**0.5 * torch.randn_like(ou)
|
||||
action = (learner.actor(history) + args.action_noise * ou).clamp(-1., 1.)
|
||||
if args.terrain_revision >= 5:
|
||||
course_types = env.scene.terrain.terrain_types.clone()
|
||||
course_levels = env.scene.terrain.terrain_levels.clone()
|
||||
course_origins = env.scene.env_origins.clone()
|
||||
prev_episode_length = env.episode_length_buf.clone()
|
||||
obs, real_reward, terminated, truncated, _ = env.step(bridge.env_action(action))
|
||||
terminated, truncated = terminated.clone(), truncated.clone()
|
||||
obs, bad = recover_invalid_envs(env, obs, real_reward)
|
||||
invalid_env_resets += int(bad.sum())
|
||||
valid = ~(terminated | truncated | bad)
|
||||
done = (terminated | truncated) & ~bad
|
||||
if done.any():
|
||||
episode_length_sum += float(prev_episode_length[done].sum().item())
|
||||
episode_count += int(done.sum())
|
||||
for name in termination_counts:
|
||||
termination_counts[name] += int(
|
||||
(env.termination_manager.get_term(name) & done).sum().item())
|
||||
real_reward = torch.where(bad, torch.zeros_like(real_reward), real_reward)
|
||||
nxt, delta = bridge.state(), pose_delta(pose, bridge.pose())
|
||||
measured = torch.cat((nxt[:, 4:6], nxt[:, 9:10]), -1)
|
||||
reward_totals += env.reward_manager._step_reward[~bad].sum(0)
|
||||
totals[0] += real_reward.sum()
|
||||
totals[1] += terminated.sum()
|
||||
totals[2] += valid.sum()
|
||||
totals[3:] += (measured[valid] - command[valid]).abs().sum(0)
|
||||
interval_steps += 1
|
||||
course_data = {}
|
||||
if args.terrain_revision >= 7:
|
||||
label_valid = env._nonwheel_label_valid & ~bad & torch.isfinite(state).all(-1)
|
||||
course_data.update(collision=env._nonwheel_hit[:,None].float(),
|
||||
collision_valid=label_valid[:,None].float())
|
||||
nonwheel_terminations += int((env._nonwheel_hit & label_valid).sum())
|
||||
if args.terrain_revision >= 5:
|
||||
new_pose = bridge.pose()
|
||||
ground_rise = field.sample(new_pose[:,:2]) - field.sample(pose[:,:2])
|
||||
forward_motion = (delta[:,:2] * command[:,:2]).sum(-1) > 0
|
||||
offset = bridge.robot.data.root_link_pos_w - course_origins
|
||||
from .stair_curriculum import stair_pass
|
||||
reached = stair_pass(offset, bridge.robot.data.projected_gravity_b,
|
||||
course_levels, course_types == 1)
|
||||
crossing = (pose[:,0] - course_origins[:,0]).abs() <= 2.8
|
||||
useful = (ground_rise > .005) & forward_motion & valid
|
||||
course_data['course'] = torch.stack((course_types.float(), course_levels.float(),
|
||||
useful.float(), (reached & crossing & valid).float()), -1)
|
||||
replay.add(state=state, action=action, previous=previous, context=context,
|
||||
next=nxt, pose=delta, valid=valid[:, None], history=history,
|
||||
command=command, bias=bias, **({"world_pose": pose} if args.terrain else {}),
|
||||
**({"next_contacts":bridge.contacts()} if contact_targets else {}),
|
||||
**({"pair_timer": env._s10_pair_timer} if args.terrain_revision >= 7 else {}),
|
||||
**course_data)
|
||||
if valid.any():
|
||||
model_context = context if field is None else torch.cat((context, field.scan(state, pose)), -1)
|
||||
learner.model.update_stats(state[valid], action[valid], previous[valid], model_context[valid],
|
||||
nxt[valid], delta[valid])
|
||||
previous = action.clone()
|
||||
previous[~valid], ou[~valid] = 0., 0.
|
||||
learner.env_steps += 1
|
||||
collect_end = time.monotonic()
|
||||
collection_time_sum += collect_end - iter_start
|
||||
wm_loss = actor_loss = critic_loss = None
|
||||
if replay.size >= args.wm_window:
|
||||
for _ in range(args.wm_updates):
|
||||
batch = (replay.sample_world_model(args.wm_batch, args.wm_window, capacity=args.archive_size)
|
||||
if args.prioritized_replay else replay.sample(args.wm_batch, args.wm_window))
|
||||
wm_loss = learner.fit(batch)
|
||||
# Replay is not checkpointed; refit it during warmup after every resume.
|
||||
if step - start >= args.warmup:
|
||||
seed = {k: v[:, 0] for k, v in replay.sample(args.imag_batch, 1).items()}
|
||||
actor_loss, critic_loss = learner.improve(seed, constants, weights, bridge.to_env, env.step_dt)
|
||||
learning_time_sum += time.monotonic() - collect_end
|
||||
if (step + 1) % args.log_interval == 0 or step + 1 == args.iterations:
|
||||
mean_episode_length = episode_length_sum / max(episode_count, 1)
|
||||
# RewardManager._step_reward is raw*weight; the env reward also
|
||||
# multiplies by dt. Scale before turning rates into episode sums.
|
||||
weighted_terms = (reward_totals / max(interval_steps * args.num_envs, 1)).cpu()
|
||||
terms_per_step = weighted_terms * env.step_dt
|
||||
episode_reward_terms = dict(zip(
|
||||
env.reward_manager.active_terms,
|
||||
(terms_per_step * mean_episode_length).tolist()))
|
||||
termination_rates = {
|
||||
name: count / max(episode_count, 1)
|
||||
for name, count in termination_counts.items()}
|
||||
now = time.monotonic()
|
||||
row = dict(step=step + 1, actor_lr=learner.optimizers[1].param_groups[0]["lr"],
|
||||
archive_windows=0 if replay.archive is None else replay.archive.size, env_steps=learner.env_steps,
|
||||
total_transitions=learner.env_steps * args.num_envs,
|
||||
num_steps_per_env=args.num_steps_per_env, wm_loss=wm_loss, actor_loss=actor_loss,
|
||||
critic_loss=critic_loss, reward=(totals[0] / (interval_steps * args.num_envs)).item(),
|
||||
falls=int(totals[1]), interval_steps=interval_steps,
|
||||
velocity_mae=(totals[3:] / totals[2].clamp_min(1)).tolist(),
|
||||
mean_episode_length=mean_episode_length,
|
||||
episode_reward_terms=episode_reward_terms,
|
||||
weighted_reward_terms=dict(zip(env.reward_manager.active_terms,
|
||||
weighted_terms.tolist())),
|
||||
episode_termination_rates=termination_rates,
|
||||
collection_time_s=collection_time_sum / max(interval_steps, 1),
|
||||
learning_time_s=learning_time_sum / max(interval_steps, 1),
|
||||
iteration_time_s=(now - began) / max(step + 1 - start, 1),
|
||||
policy_updates=learner.policy_updates,
|
||||
skipped_updates=learner.skipped_updates, invalid_env_resets=invalid_env_resets,
|
||||
elapsed_s=now - began)
|
||||
if args.task == "s10-v4-upward":
|
||||
stats = upward_curriculum_stats(env)
|
||||
row.update(stats)
|
||||
row.update(upward_tilt_stage=stats["tilt_level_q75"],
|
||||
upward_tilt_deg=TILT_STAGES_DEG[stats["tilt_level_q75"]],
|
||||
upward_recovery_rate=stats["recovery_window_rate"],
|
||||
upward_traverse_rate=stats["traverse_window_rate"],
|
||||
upward_success_rate=stats["traverse_window_rate"])
|
||||
if args.terrain_revision >= 7:
|
||||
row['nonwheel_terminations'] = nonwheel_terminations
|
||||
if args.terrain:
|
||||
row.update({k:v.item() for k,v in learner.model.last_contact_metrics.items()})
|
||||
row["terrain_level_mean"] = env.scene.terrain.terrain_levels.float().mean().item()
|
||||
row["terrain_level_max"] = env.scene.terrain.terrain_levels.max().item()
|
||||
if args.terrain_revision >= 5:
|
||||
velocity_ema = [.9*a + .1*b for a,b in zip(velocity_ema, row['velocity_mae'])]
|
||||
row['terrain_details'] = summary(env)
|
||||
row['success_archive_windows'] = 0 if replay.success_archive is None else replay.success_archive.size
|
||||
promoted = command_stage
|
||||
if args.command_curriculum and args.terrain_revision >= 5:
|
||||
promoted = next_stage(command_stage, step + 1 - args.command_curriculum_start,
|
||||
env._stair_stats.cpu().numpy(), velocity_ema)
|
||||
row['next_command_stage'] = promoted
|
||||
if args.command_curriculum:
|
||||
ranges = env.command_manager.get_term("twist").cfg.ranges
|
||||
row.update(command_stage=command_stage, command_ranges=[ranges.lin_vel_x, ranges.lin_vel_y, ranges.ang_vel_z])
|
||||
else:
|
||||
ranges = env.command_manager.get_term("twist").cfg.ranges
|
||||
row.update(command_stage=0, command_mode="fixed",
|
||||
command_ranges=[ranges.lin_vel_x, ranges.lin_vel_y, ranges.ang_vel_z])
|
||||
if args.terrain_revision >= 3:
|
||||
row["command_negative_vx_fraction"] = (command[:,0] < 0).float().mean().item()
|
||||
row["command_abs_max"] = command.abs().amax(0).tolist()
|
||||
log.write(json.dumps(row) + "\n")
|
||||
log.flush()
|
||||
print(format_training_log(row, args, episode_reward_terms, termination_rates),
|
||||
flush=True)
|
||||
if (args.terrain_revision >= 5 and args.command_curriculum
|
||||
and promoted != command_stage):
|
||||
command_stage = promoted
|
||||
env._stair_stats.zero_() # Fresh performance window for the next stage.
|
||||
set_command_stage(env.command_manager.get_term('twist').cfg, 0,
|
||||
stage=command_stage,
|
||||
stages=stages_for(args.terrain_revision))
|
||||
reward_totals.zero_()
|
||||
totals.zero_()
|
||||
interval_steps = 0
|
||||
episode_length_sum = 0.0
|
||||
episode_count = 0
|
||||
for name in termination_counts:
|
||||
termination_counts[name] = 0
|
||||
collection_time_sum = 0.0
|
||||
learning_time_sum = 0.0
|
||||
if (step + 1) % args.save_interval == 0 or step + 1 == args.iterations:
|
||||
if args.terrain:
|
||||
learner.terrain_state = {"levels": env.scene.terrain.terrain_levels.cpu().clone(),
|
||||
"types": env.scene.terrain.terrain_types.cpu().clone()}
|
||||
if args.terrain_revision >= 5:
|
||||
learner.curriculum_state = dict(stage=command_stage, velocity_ema=velocity_ema,
|
||||
stats=env._stair_stats.detach().cpu().clone())
|
||||
learner.save(output / f"model_{step + 1}.pt", step + 1, bridge.to_env, env.step_dt)
|
||||
if args.video and ((step + 1) % args.video_interval == 0
|
||||
or step + 1 == args.iterations):
|
||||
try:
|
||||
ranges = env.command_manager.get_term("twist").cfg.ranges
|
||||
record_policy_video(
|
||||
learner.actor, args, output, step + 1, bridge.to_env,
|
||||
upward_tilt_limit=(upward_video_tilt_limit(env)
|
||||
if args.task == "s10-v4-upward" else None),
|
||||
command_ranges=(ranges.lin_vel_x, ranges.lin_vel_y, ranges.ang_vel_z))
|
||||
except Exception as exc:
|
||||
print(f"RECOVERY skipped video at {step + 1}: {exc}", flush=True)
|
||||
print(f"S10 training complete: {output}")
|
||||
finally:
|
||||
env.close()
|
||||
@@ -0,0 +1,336 @@
|
||||
"""V4-upward helpers: sign-aware upright reward, reward clamps and sustained fall.
|
||||
|
||||
This module is used only by the ``s10-v4-upward`` task. V4 itself keeps its
|
||||
baseline exp-upright reward and immediate 70-degree fall termination.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
from mjlab.envs import mdp
|
||||
from mjlab.managers import ManagerTermBase
|
||||
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
||||
|
||||
from .world_model import GRAVITY
|
||||
|
||||
# Raw-term caps (before weight and dt) for the recovery task.
|
||||
# Chosen from v4 random-orientation samples: p99/max were approximately
|
||||
# leg_posture .78/.86, standstill 5.37/7.90, joint_torques 4772/5016,
|
||||
# unnecessary_wheel_lift 750/826, rolling_leg_motion 2.38/4.47.
|
||||
UPWARD_REWARD_CLAMPS = {
|
||||
"leg_posture": 2.0,
|
||||
"standstill": 4.0,
|
||||
"action_rate_l2": 4.0,
|
||||
"joint_torques_l2": 5000.0,
|
||||
"dof_pos_limits": 0.5,
|
||||
"unnecessary_wheel_lift": 2.0,
|
||||
"rolling_leg_motion": 5.0,
|
||||
}
|
||||
|
||||
# IsaacLab upward: upward max 4 (raw 4 x weight 1), bad-orientation -80,
|
||||
# i.e. the bad penalty is 20x the upward maximum. Here upward max is 2
|
||||
# (raw 1 x weight 2), so the same ratio is -40.
|
||||
BAD_ORIENTATION_WEIGHT = -40.0
|
||||
BAD_ORIENTATION_XY = 0.7
|
||||
# The old max(gz, |gx|-0.7, |gy|-0.7) severity saturated for gz >= 0.3 and had
|
||||
# a local maximum near 105 deg, where the |gx|-0.7 term crossed gz. The reward
|
||||
# severity is now a function of the tilt angle only: 0 at 45 deg, 1 at 180 deg,
|
||||
# monotonic in between. The correction term steepens the near-threshold part
|
||||
# while keeping a positive slope at 180 deg and a finite derivative at 45 deg.
|
||||
BAD_ORIENTATION_START_DEG = 45.0
|
||||
BAD_ORIENTATION_CORRECTION = 2.5
|
||||
|
||||
# Per-environment recovery curriculum, IsaacLab terrain-curriculum style.
|
||||
# Each env keeps the outcomes of its last HISTORY_LENGTH episodes and advances
|
||||
# when at least PROMOTE_SUCCESSES of them recovered (16/20 = 80% > 75%).
|
||||
# There is no global batch gate: terrain and tilt levels are per environment.
|
||||
TILT_STAGES_DEG = (60.0, 135.0, 180.0)
|
||||
HISTORY_LENGTH = 20
|
||||
PROMOTE_SUCCESSES = 16
|
||||
DEMOTE_SUCCESSES = 5
|
||||
|
||||
|
||||
def upright_upward_value(gravity: torch.Tensor) -> torch.Tensor:
|
||||
"""Normalized upward reward in [0, 1].
|
||||
|
||||
``projected_gravity_b[:, 2]`` is -1 upright, 0 sideways and +1 upside down.
|
||||
The old v4 exp-upright term is sign-blind: upside-down also scores 1. This
|
||||
linear form is monotonic through the whole range and has a constant gradient
|
||||
with respect to gravity_z, which is what the recovery task needs.
|
||||
"""
|
||||
return (0.5 * (1.0 - gravity[:, 2])).clamp(0.0, 1.0)
|
||||
|
||||
|
||||
def upright_upward(env, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
|
||||
"""Environment reward term replacing v4 ``upright`` for this task."""
|
||||
gravity = env.scene[asset_cfg.name].data.projected_gravity_b
|
||||
return upright_upward_value(gravity)
|
||||
|
||||
|
||||
def bad_orientation_value(gravity: torch.Tensor) -> torch.Tensor:
|
||||
"""Monotonic tilt-based bad-orientation severity in [0, 1].
|
||||
|
||||
``projected_gravity_b[:, 2]`` is -1 upright and +1 upside down, so
|
||||
``tilt = acos(-gz)`` goes from 0 to pi. Severity is 0 below
|
||||
``BAD_ORIENTATION_START_DEG``, then increases monotonically to 1 at 180 deg:
|
||||
|
||||
x = clamp((tilt - 45 deg) / 135 deg, 0, 1)
|
||||
severity = x + 2.5 * x * (1 - x)^2
|
||||
|
||||
Unlike the previous max(gz, |gx|-0.7, |gy|-0.7) formula, this has no
|
||||
saturation plateau for deep inversion and no local maximum near 105 deg.
|
||||
The polynomial correction is used instead of a fractional power because
|
||||
``x ** (1/3)`` has an infinite derivative at x=0 (upright), which would
|
||||
produce non-finite gradients during imagination.
|
||||
"""
|
||||
gz = gravity[:, 2].clamp(-1.0 + 1e-4, 1.0 - 1e-4)
|
||||
tilt = torch.acos(-gz)
|
||||
start = math.radians(BAD_ORIENTATION_START_DEG)
|
||||
x = ((tilt - start) / (math.pi - start)).clamp(0.0, 1.0)
|
||||
return x + BAD_ORIENTATION_CORRECTION * x * (1.0 - x).square()
|
||||
|
||||
|
||||
def bad_orientation_mask(gravity: torch.Tensor) -> torch.Tensor:
|
||||
"""Binary bad set used by the 4-second termination.
|
||||
|
||||
This intentionally keeps the IsaacLab-style boundary: upside-down
|
||||
hemisphere or |gx|/|gy| beyond 0.7. The reward severity above is a smooth
|
||||
tilt-based shaping term and is not used to decide termination.
|
||||
"""
|
||||
return ((gravity[:, 2] > 0.0)
|
||||
| (gravity[:, :2].abs() > BAD_ORIENTATION_XY).any(-1))
|
||||
|
||||
|
||||
def bad_orientation_penalty(env, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
|
||||
"""Continuous per-step bad-orientation penalty used by the environment."""
|
||||
gravity = env.scene[asset_cfg.name].data.projected_gravity_b
|
||||
return bad_orientation_value(gravity)
|
||||
|
||||
|
||||
def is_upward_task(env) -> bool:
|
||||
"""Identify the task from its upright reward function, not from a new flag."""
|
||||
rewards = getattr(env.cfg, "rewards", {})
|
||||
term = rewards.get("upright") if hasattr(rewards, "get") else None
|
||||
return term is not None and getattr(term, "func", None) is upright_upward
|
||||
|
||||
|
||||
def _sample_offsets(range_dict, count: int, device):
|
||||
"""Sample x/y/z/roll/pitch/yaw offsets; missing keys are zero."""
|
||||
offsets = []
|
||||
for key in ("x", "y", "z", "roll", "pitch", "yaw"):
|
||||
lo, hi = (range_dict or {}).get(key, (0.0, 0.0))
|
||||
lo, hi = float(lo), float(hi)
|
||||
if hi <= lo:
|
||||
offsets.append(torch.full((count,), lo, device=device))
|
||||
else:
|
||||
offsets.append(torch.empty(count, device=device).uniform_(lo, hi))
|
||||
return torch.stack(offsets, -1)
|
||||
|
||||
|
||||
def reset_root_state_upward(env, env_ids, pose_range, velocity_range=None,
|
||||
asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")):
|
||||
"""Reset with a per-environment roll/pitch limit.
|
||||
|
||||
``env._upward_tilt_limits`` is a per-env tensor updated by the curriculum.
|
||||
Play/evaluation has no such tensor and falls back to the full range in
|
||||
``pose_range`` (the task sets +/-pi there).
|
||||
"""
|
||||
env_ids = mdp.resolve_env_ids(env, env_ids)
|
||||
asset = env.scene[asset_cfg.name]
|
||||
count = len(env_ids)
|
||||
pose = _sample_offsets(pose_range, count, env.device)
|
||||
limits = getattr(env, "_upward_tilt_limits", None)
|
||||
if limits is None:
|
||||
roll = pose_range.get("roll", (-math.pi, math.pi))
|
||||
limit = max(abs(float(roll[0])), abs(float(roll[1])))
|
||||
limits = torch.full((count,), limit, device=env.device)
|
||||
else:
|
||||
limits = limits[env_ids]
|
||||
pose[:, 3] = (2.0 * torch.rand(count, device=env.device) - 1.0) * limits
|
||||
pose[:, 4] = (2.0 * torch.rand(count, device=env.device) - 1.0) * limits
|
||||
default = asset.data.default_root_state[env_ids].clone()
|
||||
positions = default[:, 0:3] + pose[:, 0:3] + env.scene.env_origins[env_ids]
|
||||
delta = mdp.quat_from_euler_xyz(pose[:, 3], pose[:, 4], pose[:, 5])
|
||||
orientations = mdp.quat_mul(default[:, 3:7], delta)
|
||||
velocities = default[:, 7:13] + _sample_offsets(velocity_range, count, env.device)
|
||||
asset.write_root_link_pose_to_sim(
|
||||
torch.cat((positions, orientations), dim=-1), env_ids=env_ids)
|
||||
asset.write_root_link_velocity_to_sim(velocities, env_ids=env_ids)
|
||||
|
||||
|
||||
def init_upward_state(env):
|
||||
"""Allocate the per-env recovery ring and tilt levels once."""
|
||||
if getattr(env, "_upward_history", None) is not None:
|
||||
return
|
||||
n, device = env.num_envs, env.device
|
||||
env._upward_history = torch.zeros(
|
||||
n, HISTORY_LENGTH, dtype=torch.bool, device=device)
|
||||
env._upward_traverse_history = torch.zeros_like(env._upward_history)
|
||||
env._upward_history_cursor = torch.zeros(n, dtype=torch.long, device=device)
|
||||
env._upward_history_filled = torch.zeros(n, dtype=torch.long, device=device)
|
||||
env._upward_tilt_levels = torch.zeros(n, dtype=torch.long, device=device)
|
||||
env._upward_tilt_limit_table = torch.tensor(
|
||||
[math.radians(d) for d in TILT_STAGES_DEG], device=device)
|
||||
env._upward_tilt_limits = env._upward_tilt_limit_table[0].expand(n).clone()
|
||||
|
||||
|
||||
def upward_curriculum(env, env_ids):
|
||||
"""Per-env rolling recovery curriculum (IsaacLab terrain-curriculum style).
|
||||
|
||||
Each environment stores the outcomes of its last ``HISTORY_LENGTH``
|
||||
episodes. ``recovered`` means the episode did not end in a bad-orientation
|
||||
termination and the final projected gravity is not in the bad set. An env
|
||||
advances when at least ``PROMOTE_SUCCESSES`` of the window recovered and
|
||||
drops when at most ``DEMOTE_SUCCESSES`` did. Terrain and tilt levels are
|
||||
updated independently per environment; no global batch rate is involved.
|
||||
"""
|
||||
init_upward_state(env)
|
||||
env_ids = mdp.resolve_env_ids(env, env_ids)
|
||||
terrain = env.scene.terrain
|
||||
data = env.scene["robot"].data
|
||||
active = env.episode_length_buf[env_ids] > 0
|
||||
fallen = env.termination_manager.terminated[env_ids]
|
||||
gravity = data.projected_gravity_b[env_ids]
|
||||
offset = data.root_link_pos_w[env_ids] - env.scene.env_origins[env_ids]
|
||||
recovered = active & ~fallen & ~bad_orientation_mask(gravity)
|
||||
traversed = (active & ~fallen & (offset[:, 0].abs() > 5.4)
|
||||
& (offset[:, 1].abs() < 6.5))
|
||||
ids = env_ids[active]
|
||||
if len(ids):
|
||||
cursor = env._upward_history_cursor[ids]
|
||||
env._upward_history[ids, cursor] = recovered[active]
|
||||
env._upward_traverse_history[ids, cursor] = traversed[active]
|
||||
env._upward_history_cursor[ids] = (cursor + 1) % HISTORY_LENGTH
|
||||
env._upward_history_filled[ids] = (
|
||||
env._upward_history_filled[ids] + 1).clamp_max(HISTORY_LENGTH)
|
||||
ready = env._upward_history_filled[env_ids] >= HISTORY_LENGTH
|
||||
successes = env._upward_history[env_ids].sum(-1)
|
||||
promote = ready & (successes >= PROMOTE_SUCCESSES)
|
||||
demote = ready & (successes <= DEMOTE_SUCCESSES)
|
||||
delta = promote.long() - demote.long()
|
||||
old_levels = terrain.terrain_levels[env_ids].clone()
|
||||
levels = (old_levels + delta).clamp(0, terrain.terrain_origins.shape[0] - 1)
|
||||
old_tilt_levels = env._upward_tilt_levels[env_ids].clone()
|
||||
tilt_levels = (old_tilt_levels + delta).clamp(0, len(TILT_STAGES_DEG) - 1)
|
||||
# A changed environment must prove itself for a fresh 20-episode window at
|
||||
# the new level; otherwise a long success streak could chain promotions in
|
||||
# consecutive episodes. A clamped promote/demote at the min/max level is not
|
||||
# a change and must keep its window so the logged rate stays meaningful.
|
||||
changed = (levels != old_levels) | (tilt_levels != old_tilt_levels)
|
||||
if changed.any():
|
||||
changed_ids = env_ids[changed]
|
||||
env._upward_history[changed_ids] = False
|
||||
env._upward_traverse_history[changed_ids] = False
|
||||
env._upward_history_filled[changed_ids] = 0
|
||||
env._upward_history_cursor[changed_ids] = 0
|
||||
terrain.terrain_levels[env_ids] = levels
|
||||
terrain.env_origins[env_ids] = terrain.terrain_origins[
|
||||
levels, terrain.terrain_types[env_ids]]
|
||||
env._upward_tilt_levels[env_ids] = tilt_levels
|
||||
env._upward_tilt_limits[env_ids] = env._upward_tilt_limit_table[tilt_levels]
|
||||
return terrain.terrain_levels.float().mean()
|
||||
|
||||
|
||||
def upward_curriculum_stats(env):
|
||||
"""Summary for logging and for choosing the training-video tilt level."""
|
||||
init_upward_state(env)
|
||||
levels = env._upward_tilt_levels
|
||||
terrain_levels = env.scene.terrain.terrain_levels
|
||||
ready = env._upward_history_filled >= HISTORY_LENGTH
|
||||
recovery = env._upward_history.sum(-1)
|
||||
traverse = env._upward_traverse_history.sum(-1)
|
||||
ready_count = int(ready.sum())
|
||||
recovery_rate = (float(recovery[ready].float().mean().item()) / HISTORY_LENGTH
|
||||
if ready_count else 0.0)
|
||||
traverse_rate = (float(traverse[ready].float().mean().item()) / HISTORY_LENGTH
|
||||
if ready_count else 0.0)
|
||||
def fractions(values, count):
|
||||
return [float((values == i).float().mean().item()) for i in range(count)]
|
||||
q75 = int(torch.quantile(levels.float(), 0.75).round().item())
|
||||
return {
|
||||
"tilt_level_mean": float(levels.float().mean().item()),
|
||||
"tilt_level_max": int(levels.max().item()),
|
||||
"tilt_level_q75": q75,
|
||||
"tilt_level_frac": fractions(levels, len(TILT_STAGES_DEG)),
|
||||
"tilt_limit_mean_deg": float(
|
||||
env._upward_tilt_limits.mean().item() * 180.0 / math.pi),
|
||||
"terrain_level_frac": fractions(
|
||||
terrain_levels, int(env.scene.terrain.terrain_origins.shape[0])),
|
||||
"recovery_window_rate": recovery_rate,
|
||||
"traverse_window_rate": traverse_rate,
|
||||
"recovery_window_ready": float(ready.float().mean().item()),
|
||||
}
|
||||
|
||||
|
||||
def upward_video_tilt_limit(env):
|
||||
"""Tilt limit reached by at least 75% of environments, for the video scene."""
|
||||
q75 = upward_curriculum_stats(env)["tilt_level_q75"]
|
||||
return math.radians(TILT_STAGES_DEG[q75])
|
||||
|
||||
|
||||
def _cap(value: torch.Tensor, max_value: float) -> torch.Tensor:
|
||||
return value.clamp(max=float(max_value))
|
||||
|
||||
|
||||
def clamped_leg_posture(env, asset_cfg: SceneEntityCfg, max_value: float) -> torch.Tensor:
|
||||
from .terrain_posture import env_posture_cost
|
||||
return _cap(env_posture_cost(env, asset_cfg), max_value)
|
||||
|
||||
|
||||
def clamped_standstill(env, wheel_cfg: SceneEntityCfg, max_value: float) -> torch.Tensor:
|
||||
from s10_supermimic.env.flat import standstill_cost
|
||||
return _cap(standstill_cost(env, wheel_cfg), max_value)
|
||||
|
||||
|
||||
def clamped_action_rate(env, max_value: float) -> torch.Tensor:
|
||||
return _cap(mdp.action_rate_l2(env), max_value)
|
||||
|
||||
|
||||
def clamped_joint_torques(env, max_value: float) -> torch.Tensor:
|
||||
return _cap(mdp.joint_torques_l2(env), max_value)
|
||||
|
||||
|
||||
def clamped_dof_pos_limits(env, asset_cfg: SceneEntityCfg, max_value: float) -> torch.Tensor:
|
||||
return _cap(mdp.joint_pos_limits(env, asset_cfg), max_value)
|
||||
|
||||
|
||||
def clamped_rolling_cost(env, term: str, max_value: float) -> torch.Tensor:
|
||||
from .rolling_rewards import env_rolling_cost
|
||||
return _cap(env_rolling_cost(env, term), max_value)
|
||||
|
||||
|
||||
def apply_reward_clamps(terms: dict[str, torch.Tensor], constants: dict) -> None:
|
||||
"""Apply the same raw caps to imagined rewards used by the env task."""
|
||||
for name, max_value in constants.get("reward_clamps", {}).items():
|
||||
if name in terms:
|
||||
terms[name] = terms[name].clamp(max=float(max_value))
|
||||
|
||||
|
||||
class BadOrientationSustained(ManagerTermBase):
|
||||
"""Terminate only after bad orientation persists for ``duration_s``.
|
||||
|
||||
This is the recovery-task state machine: a random roll/pitch spawn is allowed
|
||||
to be bad for a few seconds, but the episode ends if the robot never rights
|
||||
itself. Functionally it replaces v4's immediate 70-degree ``fell_over``.
|
||||
"""
|
||||
|
||||
def __init__(self, cfg, env):
|
||||
super().__init__(env)
|
||||
self.steps = torch.zeros(env.num_envs, dtype=torch.long, device=env.device)
|
||||
|
||||
def __call__(self, env, duration_s: float, asset_cfg: SceneEntityCfg) -> torch.Tensor:
|
||||
gravity = env.scene[asset_cfg.name].data.projected_gravity_b
|
||||
# Termination keeps the binary IsaacLab-style bad set; the reward severity
|
||||
# above is only a smooth shaping term.
|
||||
bad = bad_orientation_mask(gravity)
|
||||
self.steps = torch.where(bad, self.steps + 1, torch.zeros_like(self.steps))
|
||||
duration_steps = math.ceil(duration_s / env.step_dt)
|
||||
return self.steps >= duration_steps
|
||||
|
||||
def reset(self, env_ids=None) -> None:
|
||||
if env_ids is None:
|
||||
self.steps.zero_()
|
||||
else:
|
||||
self.steps[env_ids] = 0
|
||||
@@ -0,0 +1,101 @@
|
||||
"""Always-on wheel/body projection constraints for the safe v7 stance.
|
||||
|
||||
The body is 0.60 m long in ``base_link`` (x in [-0.30, 0.30]).
|
||||
|
||||
``wheel_body_intrusion`` keeps every individual wheel out of the opposite
|
||||
three quarters: a front wheel inside the rear three quarters (x < +0.15) or a
|
||||
rear wheel inside the front three quarters (x > -0.15) shortens the support
|
||||
polygon and makes the robot tip-prone. Its hinge is exactly zero outside the
|
||||
outer quarters and grows cubically inside, so a normal extended stance is
|
||||
untouched while a wheel tucked under the body is penalised quickly.
|
||||
|
||||
``pair_body_intrusion`` adds the wider five-sixths restriction requested for
|
||||
v7: both front wheels may not sit inside the rear five sixths (x < +0.20) at
|
||||
the same time, and both rear wheels may not sit inside the front five sixths
|
||||
(x > -0.20). A single leg crossing the line is allowed; a leg only counts as
|
||||
inside after it has stayed inside continuously for ``PAIR_GRACE_S`` (0.5 s).
|
||||
The pair penalty is therefore zero during a brief swing-through and applies
|
||||
immediately once both same-side legs are armed.
|
||||
|
||||
Both terms are kinematic and need no terrain scan. Wheels are ordered
|
||||
fl, fr, hl, hr.
|
||||
"""
|
||||
import torch
|
||||
|
||||
from .rolling_rewards import wheel_geometry
|
||||
|
||||
BODY_HALF_LENGTH = 0.30
|
||||
# Each wheel must stay in the outer quarter of the body length: x > +0.15 for
|
||||
# front wheels and x < -0.15 for rear wheels.
|
||||
WHEEL_STRIP_BOUNDARY = BODY_HALF_LENGTH / 2.0
|
||||
# Rear/front 5/6 of a 0.60 m body => the boundary sits at x = +-0.20.
|
||||
PAIR_BOUNDARY = 2.0 * BODY_HALF_LENGTH / 3.0
|
||||
PAIR_GRACE_S = 0.5
|
||||
|
||||
|
||||
def wheel_body_intrusion(state, default_q, boundary=WHEEL_STRIP_BOUNDARY):
|
||||
point, _ = wheel_geometry(state, default_q)
|
||||
x = point[..., 0]
|
||||
# Front wheels must stay ahead of +L/4, rear wheels behind -L/4.
|
||||
violation = torch.cat(((boundary - x[:, :2]).clamp_min(0.),
|
||||
(x[:, 2:] + boundary).clamp_min(0.)), -1)
|
||||
normalized = violation / boundary
|
||||
return normalized.pow(3).mean(-1)
|
||||
|
||||
|
||||
def pair_region(state, default_q, boundary=PAIR_BOUNDARY):
|
||||
"""Per-wheel inside mask and normalised depth for the 5/6 restriction."""
|
||||
point, _ = wheel_geometry(state, default_q)
|
||||
x = point[..., 0]
|
||||
inside = torch.cat((x[:, :2] < boundary, x[:, 2:] > -boundary), -1)
|
||||
depth = torch.cat(((boundary - x[:, :2]).clamp_min(0.),
|
||||
(x[:, 2:] + boundary).clamp_min(0.)), -1) / boundary
|
||||
return inside, depth
|
||||
|
||||
|
||||
def pair_body_intrusion(state, default_q, timers, dt, grace=PAIR_GRACE_S):
|
||||
"""Penalise both front or both rear wheels staying inside the 5/6 region.
|
||||
|
||||
``timers`` holds the continuous inside time of each wheel. Returning it lets
|
||||
the caller keep the 0.5 s grace across environment steps. A wheel outside
|
||||
the boundary resets its own timer, so a single-leg pass-through never
|
||||
contributes.
|
||||
"""
|
||||
inside, depth = pair_region(state, default_q)
|
||||
timers = torch.where(inside, timers + dt, torch.zeros_like(timers))
|
||||
armed = timers >= grace - 1e-6
|
||||
front = armed[:, 0] & armed[:, 1]
|
||||
rear = armed[:, 2] & armed[:, 3]
|
||||
cost = (front.float() * depth[:, :2].mean(-1)
|
||||
+ rear.float() * depth[:, 2:].mean(-1))
|
||||
return cost, timers
|
||||
|
||||
|
||||
def _bridge(env):
|
||||
from .terrain import TerrainBridge, TerrainField
|
||||
bridge = getattr(env, '_s10_posture_bridge', None)
|
||||
if bridge is None:
|
||||
bridge = TerrainBridge(env, TerrainField.from_env(env))
|
||||
env._s10_posture_bridge = bridge
|
||||
return bridge
|
||||
|
||||
|
||||
def env_wheel_body_intrusion(env):
|
||||
bridge = _bridge(env)
|
||||
return wheel_body_intrusion(
|
||||
bridge.state(), bridge.robot.data.default_joint_pos[0, bridge.legs])
|
||||
|
||||
|
||||
def env_pair_body_intrusion(env):
|
||||
bridge = _bridge(env)
|
||||
timers = getattr(env, '_s10_pair_timer', None)
|
||||
if timers is None or timers.shape[0] != env.num_envs:
|
||||
timers = torch.zeros(env.num_envs, 4, device=env.device)
|
||||
# episode_length_buf is 1 on the first reward step after an auto-reset.
|
||||
timers = timers.clone()
|
||||
timers[env.episode_length_buf <= 1] = 0.
|
||||
cost, timers = pair_body_intrusion(
|
||||
bridge.state(), bridge.robot.data.default_joint_pos[0, bridge.legs],
|
||||
timers, env.step_dt)
|
||||
env._s10_pair_timer = timers
|
||||
return cost
|
||||
@@ -0,0 +1,81 @@
|
||||
"""V6 wheel-local terrain constraints; privileged scans never enter the actor.
|
||||
|
||||
Clearance is a geometric surrogate, not a measured persistent-contact detector.
|
||||
S10 joint geometry comes from ``rolling_rewards.wheel_geometry``.
|
||||
"""
|
||||
import torch
|
||||
from .world_model import GRAVITY, LEG_POS, LEG_VEL
|
||||
from .rolling_rewards import wheel_geometry
|
||||
|
||||
|
||||
def yaw_frame_points(state, points):
|
||||
# Remove roll/pitch, retaining the heading frame used by TerrainField.scan.
|
||||
up = -torch.nn.functional.normalize(state[:, GRAVITY], dim=-1, eps=1e-6)
|
||||
cp = (up[:, 1:].square().sum(-1)+1e-12).sqrt()
|
||||
sr, cr, sp = up[:, 1] / cp, up[:, 2] / cp, -up[:, 0]
|
||||
x, y, z = points.unbind(-1)
|
||||
yr = cr[:, None]*y - sr[:, None]*z
|
||||
zr = sr[:, None]*y + cr[:, None]*z
|
||||
return torch.stack((cp[:, None]*x + sp[:, None]*zr, yr,
|
||||
-sp[:, None]*x + cp[:, None]*zr), -1)
|
||||
|
||||
|
||||
def sample_scan(scan, xy):
|
||||
"""Bilinear samples on the existing 21x13, 10cm heading-frame grid."""
|
||||
shape = xy.shape[:-1]
|
||||
x = ((xy[..., 0]+1.)*10).clamp(0,20).reshape(len(scan),-1)
|
||||
y = ((xy[..., 1]+.6)*10).clamp(0,12).reshape(len(scan),-1)
|
||||
ix, iy = x.floor().long().clamp(0,19), y.floor().long().clamp(0,11)
|
||||
ax, ay = x-ix, y-iy
|
||||
def at(dx,dy): return scan.gather(1,(ix+dx)*13+iy+dy)
|
||||
return (((1-ax)*at(0,0)+ax*at(1,0))*(1-ay)
|
||||
+ ((1-ax)*at(0,1)+ax*at(1,1))*ay).reshape(shape)
|
||||
|
||||
|
||||
def wheel_terms(state, command, scan, default_q):
|
||||
point, radius = wheel_geometry(state, default_q)
|
||||
point = yaw_frame_points(state, point)
|
||||
xy = point[..., :2]
|
||||
offsets = state.new_tensor([[x,y] for x in (-.15,0.,.15) for y in (-.10,0.,.10)])
|
||||
patch_xy = xy[:,:,None,:]+offsets
|
||||
samples = sample_scan(scan, patch_xy)
|
||||
# Index-based extrema avoid division by the number of equal extrema in
|
||||
# compiled backward through interpolated, nearly-flat patches.
|
||||
relief = samples.max(dim=-1).values-samples.min(dim=-1).values
|
||||
# Outside the scan, never assume flatness and penalize a recovering leg.
|
||||
inside = ((patch_xy[...,0].abs()<=1.) & (patch_xy[...,1].abs()<=.6)).all(-1)
|
||||
flat = (1.-(relief-.02)/.04).clamp(0,1)*inside
|
||||
straight = torch.exp(-(command[:,1]/.25).square()-(command[:,2]).square())
|
||||
gate = flat*straight[:,None]
|
||||
below = sample_scan(scan,xy)
|
||||
bottom = point[...,2]+below-radius
|
||||
lift = ((bottom-.005).clamp_min(0)/.05).square()
|
||||
tilt = ((state[:,GRAVITY][:,:2].square().sum(-1)+1e-12).sqrt()-.08).div(.22).clamp(0,1)
|
||||
motion = state[:,LEG_VEL].reshape(-1,4,3).square().mean(-1)*(1.-.8*tilt[:,None])
|
||||
q = state[:,LEG_POS].reshape(-1,4,3).square()
|
||||
posture_multiplier = torch.where(inside, 1.-.875*(relief/.12).clamp(0,1), .125)
|
||||
posture = 2.*q[:,:,0].sum(-1)+(posture_multiplier*q[:,:,1:].sum(-1)).sum(-1)
|
||||
# Only approaching a higher surface: no reward for lifting on flat ground.
|
||||
speed = command[:,:2].norm(dim=-1)
|
||||
heading_velocity = yaw_frame_points(state, torch.cat((command[:,:2], torch.zeros_like(command[:,:1])), -1)[:,None,:])[:,0,:2]
|
||||
direction = heading_velocity/(heading_velocity.square().sum(-1,keepdim=True)+1e-12).sqrt().clamp_min(.1)
|
||||
ahead_xy = xy+.16*direction[:,None,:]
|
||||
ahead = sample_scan(scan,ahead_xy)
|
||||
rise = below-ahead
|
||||
approaching = ((rise-.04)/.08).clamp(0,1)*(speed>.1)[:,None]*inside
|
||||
deficit = (radius+.015-point[...,2]-ahead).clamp_min(0)/.20
|
||||
clearance = (approaching*deficit/(1.+deficit)).mean(-1)
|
||||
return dict(unnecessary_wheel_lift=(gate*lift).sum(-1),
|
||||
rolling_leg_motion=(gate*motion).mean(-1), leg_posture=posture,
|
||||
wheel_obstacle_clearance=clearance)
|
||||
|
||||
|
||||
def env_wheel_term(env, term):
|
||||
from .terrain import TerrainBridge,TerrainField
|
||||
bridge=getattr(env,'_s10_posture_bridge',None)
|
||||
if bridge is None:
|
||||
bridge=TerrainBridge(env,TerrainField.from_env(env));env._s10_posture_bridge=bridge
|
||||
state=bridge.state()
|
||||
return wheel_terms(state,env.command_manager.get_command('twist'),
|
||||
bridge.field.scan(state,bridge.pose()),
|
||||
bridge.robot.data.default_joint_pos[0,bridge.legs])[term]
|
||||
@@ -0,0 +1,118 @@
|
||||
"""S10 residual dynamics with position legs and velocity wheels.
|
||||
|
||||
Canonical order: fl, fr, hl, hr; hipx/hipy/knee per leg, then four wheels.
|
||||
Actions are normalized policy outputs: leg targets = default + scale*a[:12],
|
||||
wheel targets = clamp(35*a[12:], -35, 35) rad/s. Wheel phase is excluded.
|
||||
This module is self-contained and depends only on torch.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
STATE_DIM = 38
|
||||
ACTION_DIM = 16
|
||||
CONTEXT_DIM = 7 # four wheel sliding frictions, base COM xyz offset (metres)
|
||||
POSE_DIM = 3
|
||||
Z, GRAVITY, LIN_VEL, ANG_VEL = slice(0, 1), slice(1, 4), slice(4, 7), slice(7, 10)
|
||||
LEG_POS, LEG_VEL, WHEEL_VEL = slice(10, 22), slice(22, 34), slice(34, 38)
|
||||
FORMAT = "s10-flat-dynamics-v1"
|
||||
|
||||
|
||||
class RunningNorm(nn.Module):
|
||||
def __init__(self, dim):
|
||||
super().__init__()
|
||||
self.register_buffer("mean", torch.zeros(dim))
|
||||
self.register_buffer("var", torch.ones(dim))
|
||||
self.register_buffer("count", torch.tensor(1e-4))
|
||||
|
||||
@torch.no_grad()
|
||||
def update(self, x):
|
||||
x = x.reshape(-1, self.mean.numel())
|
||||
if not len(x):
|
||||
return
|
||||
delta = x.mean(0) - self.mean
|
||||
total = self.count + len(x)
|
||||
self.var.copy_((self.var * self.count + x.var(0, unbiased=False) * len(x)
|
||||
+ delta.square() * self.count * len(x) / total) / total)
|
||||
self.mean.add_(delta * len(x) / total)
|
||||
self.count.copy_(total)
|
||||
|
||||
@property
|
||||
def scale(self):
|
||||
return self.var.clamp_min(1e-6).sqrt()
|
||||
|
||||
def forward(self, x):
|
||||
return ((x - self.mean) / self.scale).clamp(-10, 10)
|
||||
|
||||
|
||||
class S10WorldModel(nn.Module):
|
||||
"""(state38, action16, previous16, context7) -> (next state38, pose delta3).
|
||||
|
||||
Flat ground only. Linear and angular velocity use the base frame. Planar
|
||||
displacement uses the previous yaw frame. No target/chase or terrain scan.
|
||||
Policy history is separate from physical state and is not fed to this model.
|
||||
"""
|
||||
def __init__(self, hidden=(512, 512, 512)):
|
||||
super().__init__()
|
||||
self.hidden = tuple(hidden)
|
||||
dims = (STATE_DIM + 2 * ACTION_DIM + CONTEXT_DIM, *hidden, STATE_DIM + POSE_DIM)
|
||||
layers = []
|
||||
for i, (a, b) in enumerate(zip(dims[:-1], dims[1:])):
|
||||
layers.append(nn.Linear(a, b))
|
||||
if i < len(dims) - 2:
|
||||
layers.append(nn.SiLU())
|
||||
self.net = nn.Sequential(*layers)
|
||||
self.input_norm = RunningNorm(dims[0])
|
||||
self.output_norm = RunningNorm(dims[-1])
|
||||
|
||||
@staticmethod
|
||||
def inputs(state, action, previous_action, context):
|
||||
for name, x, dim in (("state", state, STATE_DIM), ("action", action, ACTION_DIM),
|
||||
("previous_action", previous_action, ACTION_DIM),
|
||||
("context", context, CONTEXT_DIM)):
|
||||
if x.ndim != 2 or x.shape != (state.shape[0], dim):
|
||||
raise ValueError(f"{name}: expected (batch, {dim}), got {tuple(x.shape)}")
|
||||
return torch.cat((state, action, previous_action, context), -1)
|
||||
|
||||
@torch.no_grad()
|
||||
def update_stats(self, state, action, previous_action, context, next_state, pose_delta):
|
||||
self.input_norm.update(self.inputs(state, action, previous_action, context))
|
||||
self.output_norm.update(torch.cat((next_state - state, pose_delta), -1))
|
||||
|
||||
def forward(self, state, action, previous_action, context):
|
||||
x = self.inputs(state, action, previous_action, context)
|
||||
delta = self.net(self.input_norm(x)) * self.output_norm.scale + self.output_norm.mean
|
||||
return state + delta[:, :STATE_DIM], delta[:, STATE_DIM:]
|
||||
|
||||
def loss(self, state, actions, previous_actions, context, targets, poses, valid):
|
||||
"""Autoregressive window loss; validity must be false at/reset after done.
|
||||
|
||||
Tensor shapes: state[B,38], actions/previous[B,K,16], context[B,7],
|
||||
targets[B,K,38], poses[B,K,3], valid[B,K]. A reset invalidates the rest
|
||||
of that window, even if later transitions are individually valid.
|
||||
"""
|
||||
mask = valid.bool().cumprod(1).bool()
|
||||
total = next(self.parameters()).sum() * 0.0
|
||||
for t in range(actions.shape[1]):
|
||||
# Do not execute invalid rows: NaN terminal/reset data must not leak
|
||||
# into either the loss or gradients through multiplication by zero.
|
||||
keep = mask[:, t]
|
||||
if not keep.any():
|
||||
break
|
||||
pred, pose = self(state[keep], actions[keep, t], previous_actions[keep, t], context[keep])
|
||||
error = torch.cat((pred - targets[keep, t], pose - poses[keep, t]), -1)
|
||||
total = total + (error / self.output_norm.scale).square().mean(-1).sum()
|
||||
state = state.clone()
|
||||
state[keep] = pred
|
||||
return total / mask.sum().clamp_min(1)
|
||||
|
||||
def checkpoint(self):
|
||||
return {"format": FORMAT, "hidden": self.hidden, "model": self.state_dict()}
|
||||
|
||||
@classmethod
|
||||
def from_checkpoint(cls, checkpoint, device="cpu"):
|
||||
if checkpoint.get("format") != FORMAT:
|
||||
raise ValueError("Expected a S10 dynamics checkpoint; other-robot checkpoints are incompatible")
|
||||
model = cls(checkpoint["hidden"]).to(device)
|
||||
model.load_state_dict(checkpoint["model"])
|
||||
return model
|
||||
@@ -0,0 +1,296 @@
|
||||
"""World-model future preview for the play scripts.
|
||||
|
||||
The actor checkpoint already contains the world model (``world_model``) and the
|
||||
training config. We use it directly; no separate ONNX export is required.
|
||||
|
||||
The preview has two independent outputs:
|
||||
|
||||
* a kinematic ghost trail: every predicted state is mapped back to MJCF qpos
|
||||
and transparently added to the offscreen renderer;
|
||||
* a prediction-error ledger: predictions are stored with their target step and
|
||||
compared with the real state when that step arrives (xy/z/yaw/gravity/leg q).
|
||||
|
||||
The ghost is only a visualization. It never writes into ``env.sim.data`` or the
|
||||
physics model.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
import math
|
||||
|
||||
import mujoco
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from .imagine import frame, push_history
|
||||
from .policy import Actor
|
||||
from .terrain import advance_pose
|
||||
from .terrain_world_model import TerrainWorldModel
|
||||
from .world_model import GRAVITY, LEG_POS, S10WorldModel
|
||||
|
||||
DEFAULT_HORIZON = 32
|
||||
DEFAULT_STRIDE = 4
|
||||
GHOST_NEAR = (1.0, 0.86, 0.08, 0.20)
|
||||
GHOST_FAR = (1.0, 0.24, 0.02, 0.055)
|
||||
|
||||
|
||||
@dataclass
|
||||
class GhostPrediction:
|
||||
horizon: int
|
||||
state: np.ndarray # 38
|
||||
pose: np.ndarray # x, y, yaw
|
||||
qpos: np.ndarray # MJCF qpos
|
||||
|
||||
|
||||
def load_world_model(checkpoint, device="cpu"):
|
||||
payload = checkpoint.get("world_model")
|
||||
if not isinstance(payload, dict):
|
||||
raise ValueError("Checkpoint has no world_model payload")
|
||||
fmt = payload.get("format")
|
||||
if fmt == "s10-terrain-dynamics-v1":
|
||||
model = TerrainWorldModel.from_checkpoint(payload, device)
|
||||
elif fmt == "s10-flat-dynamics-v1":
|
||||
model = S10WorldModel.from_checkpoint(payload, device)
|
||||
else:
|
||||
raise ValueError(f"Unsupported world-model format: {fmt}")
|
||||
return model.to(device).eval()
|
||||
|
||||
|
||||
def _quat_mul(a, b):
|
||||
w1, x1, y1, z1 = a
|
||||
w2, x2, y2, z2 = b
|
||||
return np.array([
|
||||
w1*w2 - x1*x2 - y1*y2 - z1*z2,
|
||||
w1*x2 + x1*w2 + y1*z2 - z1*y2,
|
||||
w1*y2 - x1*z2 + y1*w2 + z1*x2,
|
||||
w1*z2 + x1*y2 - y1*x2 + z1*w2,
|
||||
], dtype=np.float64)
|
||||
|
||||
|
||||
def _quat_from_up(up_b, yaw):
|
||||
"""Rotation mapping body-frame up to world +z, followed by world yaw."""
|
||||
up = np.asarray(up_b, dtype=np.float64)
|
||||
up /= max(np.linalg.norm(up), 1e-9)
|
||||
target = np.array([0., 0., 1.])
|
||||
dot = float(np.clip(up @ target, -1., 1.))
|
||||
if dot < -0.999999:
|
||||
align = np.array([0., 1., 0., 0.])
|
||||
else:
|
||||
axis = np.cross(up, target)
|
||||
align = np.concatenate(([1. + dot], axis))
|
||||
align /= max(np.linalg.norm(align), 1e-9)
|
||||
yaw_q = np.array([math.cos(yaw / 2), 0., 0., math.sin(yaw / 2)])
|
||||
return _quat_mul(yaw_q, align)
|
||||
|
||||
|
||||
class WorldModelPreview:
|
||||
def __init__(self, checkpoint, env, bridge, field, actor, horizon=DEFAULT_HORIZON,
|
||||
device="cpu", env_action_order=None):
|
||||
self.env = env
|
||||
self.bridge = bridge
|
||||
self.field = field
|
||||
self.actor = actor
|
||||
self.model = load_world_model(checkpoint, device)
|
||||
self.horizon = int(horizon)
|
||||
self.device = device
|
||||
self.env_action_order = list(env_action_order or checkpoint["env_action_order"])
|
||||
self.default_q = bridge.robot.data.default_joint_pos[0, bridge.legs].detach().cpu().numpy()
|
||||
model = env.sim.mj_model
|
||||
free = np.flatnonzero(model.jnt_type == mujoco.mjtJoint.mjJNT_FREE)
|
||||
if len(free) != 1:
|
||||
raise ValueError(f"Expected one free root joint, found {len(free)}")
|
||||
self.root_qadr = int(model.jnt_qposadr[int(free[0])])
|
||||
# ``bridge.legs``/``bridge.wheels`` are entity joint indices, not MuJoCo
|
||||
# joint ids (the free joint is not in the entity list). Use the entity's
|
||||
# own qpos-address table so the ghost writes the intended joints.
|
||||
qadr = bridge.robot.indexing.joint_q_adr.cpu().numpy()
|
||||
self.leg_qadr = [int(qadr[int(j)]) for j in bridge.legs]
|
||||
self.wheel_qadr = [int(qadr[int(j)]) for j in bridge.wheels]
|
||||
self.qpos_template = env.sim.data.qpos[0].detach().cpu().numpy().copy()
|
||||
# root_link_pos_w may be env-local while qpos is world. Store the constant
|
||||
# offset so predicted future positions land in the renderer's world frame.
|
||||
self.world_offset = (self.qpos_template[self.root_qadr:self.root_qadr + 2]
|
||||
- bridge.pose()[0, :2].detach().cpu().numpy())
|
||||
self.tracker = PredictionTracker(self.horizon)
|
||||
|
||||
def _ghost(self, state, pose, horizon):
|
||||
s = state[0].detach().cpu().numpy().astype(np.float64)
|
||||
p = pose[0].detach().cpu().numpy().astype(np.float64)
|
||||
qpos = self.qpos_template.copy()
|
||||
z = float(s[0])
|
||||
if self.field is not None:
|
||||
z += float(self.field.sample(pose[:, :2])[0].detach().cpu())
|
||||
quat = _quat_from_up(-s[1:4], float(p[2]))
|
||||
qpos[self.root_qadr:self.root_qadr + 3] = (
|
||||
p[0] + self.world_offset[0], p[1] + self.world_offset[1], z)
|
||||
qpos[self.root_qadr + 3:self.root_qadr + 7] = quat
|
||||
qpos[self.leg_qadr] = s[LEG_POS] + self.default_q
|
||||
qpos[self.wheel_qadr] = self.qpos_template[self.wheel_qadr]
|
||||
return GhostPrediction(horizon, s, p, qpos)
|
||||
|
||||
@torch.no_grad()
|
||||
def predict(self, history, previous, command, context, bias, state, pose):
|
||||
"""Roll the actor + world model forward from the current real state."""
|
||||
ghosts = []
|
||||
history = history.clone()
|
||||
previous = previous.clone()
|
||||
state = state.clone()
|
||||
pose = pose.clone()
|
||||
scan = None if self.field is None else self.field.scan(state, pose)
|
||||
for i in range(self.horizon):
|
||||
action = self.actor(history)
|
||||
model_context = context if self.field is None else torch.cat((context, scan), -1)
|
||||
if self.field is not None:
|
||||
nxt, delta, _, _ = self.model.forward_all(state, action, previous, model_context)
|
||||
else:
|
||||
nxt, delta = self.model(state, action, previous, model_context)
|
||||
pose = advance_pose(pose, delta)
|
||||
if self.field is not None:
|
||||
scan = self.field.scan(nxt, pose)
|
||||
latest = frame(nxt, action[:, self.env_action_order], command, bias,
|
||||
noise=False, wheel_vel=self.actor.actor_wheel_vel)
|
||||
history = push_history(history, latest, self.actor.history, self.actor.terms)
|
||||
previous = action
|
||||
state = nxt
|
||||
ghosts.append(self._ghost(state, pose, i + 1))
|
||||
return ghosts
|
||||
|
||||
def update_errors(self, step, state, pose):
|
||||
self.tracker.update(step, state[0].detach().cpu().numpy(),
|
||||
pose[0].detach().cpu().numpy())
|
||||
|
||||
def push_errors(self, step, ghosts):
|
||||
self.tracker.push(step, ghosts)
|
||||
|
||||
def summary(self):
|
||||
return self.tracker.summary()
|
||||
|
||||
def clear(self):
|
||||
self.tracker.clear()
|
||||
|
||||
def render(self, ghosts, stride=DEFAULT_STRIDE, rgba=None):
|
||||
"""Actual frame plus transparent ghost geoms, using the env's renderer."""
|
||||
return render_with_ghosts(self.env, ghosts, stride=stride, rgba=rgba)
|
||||
|
||||
|
||||
class PredictionTracker:
|
||||
"""Compare predictions with the state that actually arrives later."""
|
||||
def __init__(self, horizon, keep=200):
|
||||
self.horizon = int(horizon)
|
||||
self.keep = int(keep)
|
||||
self.pending = []
|
||||
self.errors = {h: [] for h in (1, 4, 8, 16, 32) if h <= self.horizon}
|
||||
|
||||
def clear(self):
|
||||
self.pending.clear()
|
||||
for values in self.errors.values():
|
||||
values.clear()
|
||||
|
||||
@staticmethod
|
||||
def _error(pred_state, pred_pose, state, pose):
|
||||
dxy = float(np.linalg.norm(pred_pose[:2] - pose[:2]))
|
||||
dz = float(abs(pred_state[0] - state[0]))
|
||||
dyaw = float(abs(math.atan2(math.sin(pred_pose[2] - pose[2]),
|
||||
math.cos(pred_pose[2] - pose[2]))))
|
||||
pg = pred_state[GRAVITY]
|
||||
ag = state[GRAVITY]
|
||||
pg = pg / max(np.linalg.norm(pg), 1e-9)
|
||||
ag = ag / max(np.linalg.norm(ag), 1e-9)
|
||||
dg = float(math.degrees(math.acos(float(np.clip(pg @ ag, -1., 1.)))))
|
||||
dq = float(np.sqrt(np.mean((pred_state[LEG_POS] - state[LEG_POS]) ** 2)))
|
||||
return dxy, dz, dyaw, dg, dq
|
||||
|
||||
def push(self, step, ghosts):
|
||||
for ghost in ghosts:
|
||||
if ghost.horizon in self.errors:
|
||||
self.pending.append((step + ghost.horizon, ghost.horizon,
|
||||
ghost.state, ghost.pose))
|
||||
|
||||
def update(self, step, state, pose):
|
||||
keep = []
|
||||
for target, horizon, pred_state, pred_pose in self.pending:
|
||||
if target == step:
|
||||
values = self.errors[horizon]
|
||||
values.append(self._error(pred_state, pred_pose, state, pose))
|
||||
if len(values) > self.keep:
|
||||
del values[:-self.keep]
|
||||
elif target > step:
|
||||
keep.append((target, horizon, pred_state, pred_pose))
|
||||
self.pending = keep
|
||||
|
||||
def summary(self):
|
||||
parts = []
|
||||
for horizon in sorted(self.errors):
|
||||
values = self.errors[horizon]
|
||||
if not values:
|
||||
continue
|
||||
mean = np.mean(values, axis=0)
|
||||
parts.append(f"h{horizon} xy {mean[0]:.3f}m yaw {math.degrees(mean[2]):.1f}deg "
|
||||
f"leg {mean[4]:.3f}")
|
||||
return " | ".join(parts) if parts else "collecting..."
|
||||
|
||||
|
||||
def annotate_rgb(rgb, lines, color=(255, 185, 90)):
|
||||
"""Small bottom-left text overlay for MP4 frames (Pillow optional)."""
|
||||
if not lines:
|
||||
return rgb
|
||||
try:
|
||||
from PIL import Image, ImageDraw, ImageFont
|
||||
except ImportError:
|
||||
return rgb
|
||||
image = Image.fromarray(rgb)
|
||||
draw = ImageDraw.Draw(image)
|
||||
font = ImageFont.load_default()
|
||||
y = image.height - 8 - 14 * len(lines)
|
||||
for line in lines:
|
||||
draw.rectangle((6, y - 2, 6 + 8 * len(line), y + 13), fill=(8, 10, 14))
|
||||
draw.text((10, y), line, fill=color, font=font)
|
||||
y += 14
|
||||
return np.asarray(image)
|
||||
|
||||
|
||||
def render_with_ghosts(env, ghosts, stride=DEFAULT_STRIDE, rgba=None):
|
||||
"""Add ghost robot geoms to the existing offscreen scene and render once."""
|
||||
rgb = env.render()
|
||||
renderer = env._offline_renderer
|
||||
if renderer is None or not ghosts:
|
||||
return rgb
|
||||
import mujoco
|
||||
|
||||
model, data = renderer._model, renderer._data
|
||||
scene = renderer.renderer.scene
|
||||
old_qpos, old_qvel = data.qpos.copy(), data.qvel.copy()
|
||||
old_group = renderer._opt.geomgroup.copy()
|
||||
# S10 visual meshes are group 2; group 0 is terrain and group 1 collision
|
||||
# geometry. Adding only group 2 gives robot ghosts without duplicating terrain.
|
||||
visual = np.flatnonzero(model.geom_group == 2)
|
||||
old_geom_rgba = model.geom_rgba[visual].copy()
|
||||
mat_ids = np.unique(model.geom_matid[visual])
|
||||
mat_ids = mat_ids[mat_ids >= 0]
|
||||
old_mat_rgba = {int(i): model.mat_rgba[int(i)].copy() for i in mat_ids}
|
||||
near = np.asarray(rgba if rgba is not None else GHOST_NEAR, dtype=np.float32)
|
||||
far = np.asarray(GHOST_FAR, dtype=np.float32)
|
||||
try:
|
||||
renderer._opt.geomgroup[:] = 0
|
||||
renderer._opt.geomgroup[2] = 1
|
||||
for ghost in ghosts[max(0, int(stride) - 1)::max(1, int(stride))]:
|
||||
# Near future is yellow and stronger; the far tail fades to red.
|
||||
t = ghost.horizon / max(1, ghosts[-1].horizon)
|
||||
tint = near * (1. - t) + far * t
|
||||
for mat_id in mat_ids:
|
||||
model.mat_rgba[int(mat_id)] = tint
|
||||
model.geom_rgba[visual] = tint
|
||||
data.qpos[:] = ghost.qpos
|
||||
data.qvel[:] = 0.
|
||||
mujoco.mj_forward(model, data)
|
||||
mujoco.mjv_addGeoms(model, data, renderer._opt, renderer._pert,
|
||||
renderer._catmask.value, scene)
|
||||
finally:
|
||||
renderer._opt.geomgroup[:] = old_group
|
||||
model.geom_rgba[visual] = old_geom_rgba
|
||||
for mat_id, value in old_mat_rgba.items():
|
||||
model.mat_rgba[mat_id] = value
|
||||
data.qpos[:] = old_qpos
|
||||
data.qvel[:] = old_qvel
|
||||
mujoco.mj_forward(model, data)
|
||||
return renderer.render()
|
||||
Reference in New Issue
Block a user