783 lines
41 KiB
Python
783 lines
41 KiB
Python
"""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()
|