"""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()