microduck/microduck_rl
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Import upstream snapshot d424a0c899f6b33cbd3daeb279913134349c0b63

Browse Source 
Upstream: https://github.com/pollen-robotics/microduck_rl
Upstream-Commit: d424a0c899f6b33cbd3daeb279913134349c0b63
Upstream-Branch: develop
This commit is contained in:
Upstream Snapshot
2026-08-28 15:41:56 +08:00
commit 47372443ff
222 changed files with 38327 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/crouch_pose_editor.py
+99
View File
@@ -0,0 +1,99 @@
"""Interactive crouch pose editor (roller robot).
Ouvre le viewer MuJoCo avec le robot rollers debout. Dans le panneau "Control"
du viewer, bouge les sliders (genoux/hanches/chevilles…) pour composer la pose
ACCROUPIE voulue. La gravité est coupée et la base est maintenue droite +
abaissée pour que le point le plus bas reste au sol (tu vois donc le tronc
descendre quand tu plies les genoux). À la fermeture de la fenêtre, la pose est
imprimée en dict CROUCH_POSE {nom_articulation: angle_rad} prêt à coller.
Usage:
uv run python scripts/crouch_pose_editor.py
"""
import re
import time
import mujoco
import mujoco.viewer
from mjlab_microduck.robot.microduck_constants import (
get_walk_rollers_spec,
HOME_FRAME,
)
def home_value(joint_name: str):
for pattern, val in HOME_FRAME.joint_pos.items():
if re.search(pattern, joint_name):
return float(val)
return 0.0
# Modèle direct depuis le spec du robot (14 actionneurs <position> dans le XML).
model = get_walk_rollers_spec().compile()
data = mujoco.MjData(model)
mujoco.mj_resetData(model, data)
model.opt.gravity[:] = [0, 0, 0] # rien ne s'effondre : seuls les sliders bougent
has_free = model.jnt_type[0] == mujoco.mjtJoint.mjJNT_FREE
# Articulations actionnées (hors roues passives), avec adresse qpos.
joints = []
for i in range(model.njnt):
name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_JOINT, i)
if not name or "freejoint" in name or "passive_" in name:
continue
joints.append((name, model.jnt_qposadr[i]))
# ctrl initial = pose HOME (les actionneurs position tiennent cette cible).
for a in range(model.nu):
aname = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_ACTUATOR, a)
data.ctrl[a] = home_value(aname or "")
if has_free:
data.qpos[0:3] = [0.0, 0.0, 0.14]
data.qpos[3:7] = [1.0, 0.0, 0.0, 0.0]
base_xy = data.qpos[0:2].copy()
base_quat = data.qpos[3:7].copy()
robot_geoms = [g for g in range(model.ngeom)
if model.geom_type[g] != mujoco.mjtGeom.mjGEOM_PLANE]
mujoco.mj_forward(model, data)
print("=== Crouch Pose Editor (rollers) ===")
print(f"actionneurs: {model.nu} | base flottante: {has_free}")
print("Ouvre le panneau 'Control' du viewer et bouge les sliders pour composer")
print("la pose ACCROUPIE. Ferme la fenêtre quand c'est bon.\n")
with mujoco.viewer.launch_passive(model, data) as viewer:
while viewer.is_running():
if has_free:
data.qpos[0:2] = base_xy
data.qpos[3:7] = base_quat
data.qvel[0:6] = 0.0
mujoco.mj_step(model, data) # actionneurs position -> les joints suivent ctrl
if has_free:
data.qpos[0:2] = base_xy
data.qpos[3:7] = base_quat
data.qvel[0:6] = 0.0
mujoco.mj_forward(model, data)
try:
zmin = min(float(data.geom_xpos[g, 2] - model.geom_rbound[g])
for g in robot_geoms)
data.qpos[2] -= zmin
mujoco.mj_forward(model, data)
except Exception:
pass
viewer.sync()
time.sleep(1.0 / 60.0)
print("\n=== Pose accroupie capturée ===\n")
print("CROUCH_POSE = {")
for name, adr in joints:
print(f' "{name}": {float(data.qpos[adr]):.4f},')
print("}")
if has_free:
print(f"\n# hauteur de base finale (info) : z = {float(data.qpos[2]):.4f}")
print("# Colle CROUCH_POSE ici et donne-le a Claude pour cabler la reward.")
scripts/export.py
+279
View File
@@ -0,0 +1,279 @@
"""Script to play RL agent with RSL-RL."""
import os
import re
import sys
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Literal
import torch
import tyro
from rsl_rl.runners import OnPolicyRunner
from mjlab.envs import ManagerBasedRlEnv
from mjlab.rl import RslRlVecEnvWrapper
from mjlab.tasks.registry import list_tasks, load_env_cfg, load_rl_cfg, load_runner_cls
from mjlab.tasks.tracking.mdp import MotionCommandCfg
from mjlab.utils.os import get_checkpoint_path, get_wandb_checkpoint_path
from mjlab.utils.torch import configure_torch_backends
from mjlab.utils.wrappers import VideoRecorder
from mjlab.viewer import NativeMujocoViewer, ViserPlayViewer
@dataclass(frozen=True)
class ExportConfig:
onnx_file: str = "output.onnx"
agent: Literal["zero", "random", "trained"] = "trained"
registry_name: str | None = None
wandb_run_path: str | None = None
checkpoint: int | None = None # Select checkpoint by iteration number (e.g. 3000)
checkpoint_file: str | None = None
motion_file: str | None = None
num_envs: int | None = None
device: str | None = None
video: bool = False
video_length: int = 200
video_height: int | None = None
video_width: int | None = None
camera: int | str | None = None
viewer: Literal["auto", "native", "viser"] = "auto"
# Internal flag used by demo script.
_demo_mode: tyro.conf.Suppress[bool] = False
def run_export(task_id: str, cfg: ExportConfig):
configure_torch_backends()
device = cfg.device or ("cuda:0" if torch.cuda.is_available() else "cpu")
env_cfg = load_env_cfg(task_id, play=True)
agent_cfg = load_rl_cfg(task_id)
DUMMY_MODE = cfg.agent in {"zero", "random"}
TRAINED_MODE = not DUMMY_MODE
# Check if this is a motion tracking task.
is_motion_tracking = (
env_cfg.commands is not None
and "motion" in env_cfg.commands
and isinstance(env_cfg.commands["motion"], MotionCommandCfg)
)
is_tracking_task = is_motion_tracking
if is_tracking_task and cfg._demo_mode:
# Demo mode: use uniform sampling to see more diversity with num_envs > 1.
assert env_cfg.commands is not None
motion_cmd = env_cfg.commands["motion"]
assert isinstance(motion_cmd, MotionCommandCfg)
motion_cmd.sampling_mode = "uniform"
if is_tracking_task:
assert env_cfg.commands is not None
motion_cmd = env_cfg.commands["motion"]
assert isinstance(motion_cmd, MotionCommandCfg)
# Check if motion file is already set and exists
motion_file_already_set = (
hasattr(motion_cmd, 'motion_file')
and motion_cmd.motion_file is not None
and Path(motion_cmd.motion_file).exists()
)
if DUMMY_MODE:
if not cfg.registry_name:
raise ValueError(
"Tracking tasks require `registry_name` when using dummy agents."
)
# Check if the registry name includes alias, if not, append ":latest".
registry_name = cfg.registry_name
if ":" not in registry_name:
registry_name = registry_name + ":latest"
import wandb
api = wandb.Api()
artifact = api.artifact(registry_name)
motion_cmd.motion_file = str(Path(artifact.download()) / "motion.npz")
else:
if cfg.motion_file is not None:
print(f"[INFO]: Using motion file from CLI: {cfg.motion_file}")
motion_cmd.motion_file = cfg.motion_file
elif motion_file_already_set:
print(f"[INFO]: Using motion file from env config: {motion_cmd.motion_file}")
else:
# Try to download from wandb artifacts
import wandb
api = wandb.Api()
if cfg.wandb_run_path is None and cfg.checkpoint_file is not None:
raise ValueError(
"Tracking tasks require `motion_file` when using `checkpoint_file`, "
"or provide `wandb_run_path` so the motion artifact can be resolved."
)
if cfg.wandb_run_path is not None:
wandb_run = api.run(str(cfg.wandb_run_path))
art = next(
(a for a in wandb_run.used_artifacts() if a.type == "motions"),
None,
)
if art is None:
raise RuntimeError("No motion artifact found in the run.")
motion_cmd.motion_file = str(Path(art.download()) / "motion.npz")
log_dir: Path | None = None
resume_path: Path | None = None
if TRAINED_MODE:
log_root_path = (Path("logs") / "rsl_rl" / agent_cfg.experiment_name).resolve()
if cfg.checkpoint_file is not None:
resume_path = Path(cfg.checkpoint_file)
if not resume_path.exists():
raise FileNotFoundError(f"Checkpoint file not found: {resume_path}")
print(f"[INFO]: Loading checkpoint: {resume_path.name}")
elif cfg.checkpoint is not None:
# Select a specific checkpoint iteration, from wandb or local.
checkpoint_filename = f"model_{cfg.checkpoint}.pt"
if cfg.wandb_run_path is not None:
import wandb
api = wandb.Api()
wandb_run = api.run(str(cfg.wandb_run_path))
run_id = cfg.wandb_run_path.split("/")[-1]
download_dir = log_root_path / "wandb_checkpoints" / run_id
resume_path = download_dir / checkpoint_filename
if resume_path.exists():
print(f"[INFO]: Loading checkpoint: {checkpoint_filename} (run: {run_id}, cached)")
else:
available = [f.name for f in wandb_run.files() if "model" in f.name]
if checkpoint_filename not in available:
raise FileNotFoundError(
f"Checkpoint '{checkpoint_filename}' not found in wandb run. "
f"Available: {sorted(available)}"
)
wandb_run.file(checkpoint_filename).download(str(download_dir), replace=True)
print(f"[INFO]: Loading checkpoint: {checkpoint_filename} (run: {run_id}, downloaded)")
else:
resume_path = get_checkpoint_path(
log_root_path, checkpoint=re.escape(checkpoint_filename)
)
print(f"[INFO]: Loading checkpoint: {resume_path.name}")
else:
if cfg.wandb_run_path is None:
raise ValueError(
"`wandb_run_path` is required when `checkpoint_file` is not provided."
)
resume_path, was_cached = get_wandb_checkpoint_path(
log_root_path, Path(cfg.wandb_run_path)
)
# Extract run_id and checkpoint name from path for display.
run_id = resume_path.parent.name
checkpoint_name = resume_path.name
cached_str = "cached" if was_cached else "downloaded"
print(
f"[INFO]: Loading checkpoint: {checkpoint_name} (run: {run_id}, {cached_str})"
)
log_dir = resume_path.parent
if cfg.num_envs is not None:
env_cfg.scene.num_envs = cfg.num_envs
if cfg.video_height is not None:
env_cfg.viewer.height = cfg.video_height
if cfg.video_width is not None:
env_cfg.viewer.width = cfg.video_width
render_mode = "rgb_array" if (TRAINED_MODE and cfg.video) else None
if cfg.video and DUMMY_MODE:
print(
"[WARN] Video recording with dummy agents is disabled (no checkpoint/log_dir)."
)
env = ManagerBasedRlEnv(cfg=env_cfg, device=device, render_mode=render_mode)
if TRAINED_MODE and cfg.video:
print("[INFO] Recording videos during play")
assert log_dir is not None # log_dir is set in TRAINED_MODE block
env = VideoRecorder(
env,
video_folder=log_dir / "videos" / "play",
step_trigger=lambda step: step == 0,
video_length=cfg.video_length,
disable_logger=True,
)
env = RslRlVecEnvWrapper(env, clip_actions=agent_cfg.clip_actions)
if DUMMY_MODE:
action_shape: tuple[int, ...] = env.unwrapped.action_space.shape # type: ignore
if cfg.agent == "zero":
class PolicyZero:
def __call__(self, obs) -> torch.Tensor:
del obs
return torch.zeros(action_shape, device=env.unwrapped.device)
policy = PolicyZero()
else:
class PolicyRandom:
def __call__(self, obs) -> torch.Tensor:
del obs
return 2 * torch.rand(action_shape, device=env.unwrapped.device) - 1
policy = PolicyRandom()
else:
runner_cls = load_runner_cls(task_id) or OnPolicyRunner
runner = runner_cls(env, asdict(agent_cfg), device=device)
runner.load(str(resume_path), map_location=device)
policy = runner.get_inference_policy(device=device)
# mjlab 1.3.0: ONNX export + metadata moved to mjlab.rl.exporter_utils and
# the runner's built-in export_policy_to_onnx. Observation normalization is
# baked into the exported graph automatically — EmpiricalNormalization is a
# submodule of the policy's MLPModel (obs_normalization=True in RslRlModelCfg),
# so export_policy_to_onnx emits actor(normalizer(obs)). No manual normalizer
# handling needed (the old export_velocity_policy_as_onnx path is gone).
from mjlab.rl.exporter_utils import get_base_metadata, attach_metadata_to_onnx
onnx_path = os.path.abspath(cfg.onnx_file)
path = os.path.dirname(onnx_path)
filename = os.path.basename(onnx_path)
runner.export_policy_to_onnx(path, filename)
metadata = get_base_metadata(runner.env.unwrapped, run_path=cfg.checkpoint_file)
attach_metadata_to_onnx(onnx_path, metadata)
print(f"Written {onnx_path}")
env.close()
def main():
# Parse first argument to choose the task.
# Import tasks to populate the registry.
import mjlab.tasks # noqa: F401
all_tasks = list_tasks()
chosen_task, remaining_args = tyro.cli(
tyro.extras.literal_type_from_choices(all_tasks),
add_help=False,
return_unknown_args=True,
)
# Parse the rest of the arguments + allow overriding env_cfg and agent_cfg.
agent_cfg = load_rl_cfg(chosen_task)
args = tyro.cli(
ExportConfig,
args=remaining_args,
default=ExportConfig(),
prog=sys.argv[0] + f" {chosen_task}",
config=(
tyro.conf.AvoidSubcommands,
tyro.conf.FlagConversionOff,
),
)
del remaining_args, agent_cfg
run_export(chosen_task, args)
if __name__ == "__main__":
main()
scripts/hf/README.md
+79
View File
@@ -0,0 +1,79 @@
# HF Jobs training
Train mjlab-microduck on Hugging Face's managed GPUs. Auth is the cached HF
token (`hf auth login` or `HF_TOKEN`); everything goes through the
`huggingface_hub` Python API — the standalone `hf` CLI is not required.
## One-time setup
```fish
hf auth login # or export HF_TOKEN (any tool that caches the token works)
wandb login # auto-detected from ~/.netrc and forwarded
```
## Submit a run
Your normal train command, plus `--hf-jobs`:
```fish
uv run train Mjlab-Kick-Flat-MicroDuck \
--env.scene.num-envs 4096 --agent.max_iterations 4000 --hf-jobs
```
You'll be asked which namespace to run under — your personal account or one
of your orgs. Repos, uv-cache bucket, billing and the job itself all live in
the chosen namespace. Pass `--namespace <name>` to skip the prompt
(non-interactive runs default to personal).
Without `--hf-jobs` the command behaves exactly as before (local training).
Submission flags are consumed locally; everything else is forwarded to
`uv run train` inside the job.
Useful flags:
- `--namespace <name>` — account/org to run under; skips the prompt
- `--flavor l4x1` (default) / `a10g-large` / `a100-large`
- `--timeout 12h` (default) — job is killed past this
- `--detach` — submit and return immediately (default streams logs; Ctrl-C detaches without killing the job)
- `--dry-run` — build tarball, print the job spec, do not submit
- `--run-name <tag>` — overrides the auto-generated `<task>-<timestamp>` name
- `--no-uv-cache` — disable the persistent `uv` cache bucket (first-run cost on every run)
- `--no-wandb` — don't forward a wandb key
(`uv run scripts/hf/train_hf.py <task> ...` still works — it's a shim to the
same code, which lives in `src/mjlab_microduck/hf_jobs.py`.)
## What happens under the hood
1. `git ls-files` snapshots tracked + uncommitted files of the repo you run
from (worktree-aware) → `src-<stamp>.tar.gz`.
2. Tarball is uploaded to private dataset `<namespace>/mjlab-microduck-src`.
3. Private model repo `<namespace>/<run-name>` is created for checkpoints.
4. A private HF bucket `<namespace>/mjlab-uv-cache` is mounted at `/uv-cache`
and used as `UV_CACHE_DIR` so wheel downloads persist across runs (first
run cold, subsequent runs fast).
5. `HfApi.run_job` launches a container that:
- installs `uv`, extracts the tarball, runs `uv sync` (warm-cached),
- starts `scripts/hf/uploader.py` in background (watches `logs/rsl_rl/**/model_*.pt`, pushes every 60s),
- runs `uv run train <task> <args>`,
- does a final one-shot upload on exit.
6. wandb credentials are forwarded as a secret — runs show up live in your
wandb project.
## Browsing checkpoints
The submitter prints `https://huggingface.co/<namespace>/<run-name>` at
start; new `.pt` files appear there during training.
## Managing jobs
The job id and URL are printed at submission. From Python:
```python
from huggingface_hub import HfApi
api = HfApi()
api.list_jobs() # or namespace="pollen-robotics"
for l in api.fetch_job_logs(job_id="...", follow=True): print(l)
api.cancel_job(job_id="...")
```
(or the `hf jobs ps/logs/cancel` CLI if you have it installed.)
scripts/hf/train_hf.py
+15
View File
@@ -0,0 +1,15 @@
"""Back-compat shim: the submission logic moved to mjlab_microduck.hf_jobs.
Prefer the integrated flag:
uv run train <task> <train args...> --hf-jobs [--namespace <ns>] [...]
This script keeps the old invocation working:
uv run scripts/hf/train_hf.py <task> [submission flags] <train args...>
"""
import sys
from mjlab_microduck.hf_jobs import submit
if __name__ == "__main__":
sys.exit(submit(sys.argv[1:]))
scripts/hf/uploader.py
+75
View File
@@ -0,0 +1,75 @@
"""Checkpoint uploader run inside an HF Job.
Watches `logs/rsl_rl/**/model_*.pt` and uploads new/updated files to the
target HF Model repo. Designed to be `nohup uv run`-launched from the job
bootstrap, with auth coming from the HF_TOKEN secret injected by `hf jobs run`.
"""
from __future__ import annotations
import os
import sys
import time
from pathlib import Path
from huggingface_hub import HfApi, CommitOperationAdd
def main() -> int:
repo_id = os.environ.get("CKPT_REPO")
if not repo_id:
print("[uploader] CKPT_REPO not set, exiting", flush=True)
return 1
poll_interval = float(os.environ.get("CKPT_POLL_INTERVAL", "60"))
root = Path(os.environ.get("CKPT_ROOT", "logs/rsl_rl"))
one_shot = os.environ.get("CKPT_ONE_SHOT") == "1"
api = HfApi()
api.create_repo(repo_id, repo_type="model", private=True, exist_ok=True)
mode = "one-shot" if one_shot else f"every {poll_interval}s"
print(f"[uploader] watching {root} -> {repo_id} ({mode})", flush=True)
sent: dict[Path, float] = {}
while True:
try:
files = list(root.glob("**/model_*.pt"))
# also pick up the dumped configs once
files += [p for p in root.glob("**/params/*.yaml")]
files += [p for p in root.glob("**/params/*.json")]
to_upload: list[CommitOperationAdd] = []
for f in files:
try:
mtime = f.stat().st_mtime
except FileNotFoundError:
continue
if sent.get(f) == mtime:
continue
# use path-in-repo relative to logs/rsl_rl so the repo mirrors run dirs
rel = f.relative_to(root)
to_upload.append(
CommitOperationAdd(path_in_repo=str(rel), path_or_fileobj=str(f))
)
sent[f] = mtime
if to_upload:
msg = f"upload {len(to_upload)} file(s)"
api.create_commit(
repo_id=repo_id,
repo_type="model",
operations=to_upload,
commit_message=msg,
)
print(f"[uploader] pushed {len(to_upload)} file(s)", flush=True)
except Exception as e:
print(f"[uploader] error: {e}", flush=True)
if one_shot:
return 0
time.sleep(poll_interval)
if __name__ == "__main__":
sys.exit(main())
scripts/infer_policy.py
+1383
View File
File diff suppressed because it is too large Load Diff
scripts/play_latest.py
+57
View File
@@ -0,0 +1,57 @@
"""Find the latest wandb run for a given user and launch `uv run play`.
Par défaut : le tout dernier run (toutes tâches confondues).
Avec un flag de type, le dernier run de CE type seulement :
md-play --crouch # dernier Mjlab-RollerCrouch-...
md-play --roller # dernier Mjlab-...-Rollers
md-play --swizzle # dernier Mjlab-...-Swizzle-...
md-play --slope # dernier Mjlab-RollerSlope-...
Les arguments inconnus sont transmis tels quels à `uv run play`
(ex: md-play --crouch --action-scale 0.8).
"""
import argparse
from wandb_utils import resolve_run, run_command
# flag -> sous-chaîne recherchée dans le task_id (metadata args[0])
TYPE_SUBSTR = {
"crouch": "Crouch", # Mjlab-RollerCrouch-Flat-MicroDuck
"roller": "MicroDuck-Rollers", # Mjlab-Velocity-Flat-MicroDuck-Rollers (≠ RollerSlope/RollerCrouch)
"swizzle": "Swizzle", # Mjlab-Velocity-Swizzle-MicroDuck
"slope": "Slope", # Mjlab-RollerSlope-Flat-MicroDuck
}
def main():
parser = argparse.ArgumentParser(description="Play latest wandb run for a user")
parser.add_argument(
"--user", default="coralie",
help="Filter runs by user (matched against email, default: coralie)",
)
parser.add_argument(
"--dry-run", action="store_true",
help="Print the play command without executing it",
)
for t in TYPE_SUBSTR:
parser.add_argument(
f"--{t}", dest="type", action="store_const", const=t,
help=f"latest '{t}' run only",
)
# unknown args (e.g. --action-scale 0.8) are forwarded to `uv run play`
args, extra = parser.parse_known_args()
task_substr = TYPE_SUBSTR[args.type] if getattr(args, "type", None) else None
_, info = resolve_run(args.user, task_substr)
cmd = [
"uv", "run", "play",
info["env_name"],
"--wandb-run-path", info["run_path"],
*extra,
]
run_command(cmd, args.dry_run)
if __name__ == "__main__":
main()
scripts/plot_observations_comparison_plotly.py
+306
View File
@@ -0,0 +1,306 @@
#!/usr/bin/env python3
"""
Plot comparison between real and simulated observations using Plotly.
"""
import argparse
import pickle
import numpy as np
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from pathlib import Path
def load_observations(pkl_path: str):
"""Load observations from pickle file."""
with open(pkl_path, 'rb') as f:
data = pickle.load(f)
if isinstance(data, dict):
if 'observations' in data and 'timestamps' in data:
observations = data['observations']
timestamps = data['timestamps']
else:
raise ValueError("Dictionary must contain 'observations' and 'timestamps' keys")
elif isinstance(data, list):
if len(data) == 0:
raise ValueError("Empty data list")
if isinstance(data[0], dict) and 'timestamp' in data[0] and 'observation' in data[0]:
timestamps = [item['timestamp'] for item in data]
observations = [item['observation'] for item in data]
elif isinstance(data[0], tuple):
timestamps = [item[0] for item in data]
observations = [item[1] for item in data]
else:
observations = data
timestamps = [i * 0.02 for i in range(len(observations))]
else:
raise ValueError(f"Unsupported data format: {type(data)}")
return np.array(observations), np.array(timestamps)
def plot_comparison(real_obs, real_ts, sim_obs=None, sim_ts=None):
"""
Plot comparison between real and simulated observations using Plotly.
If sim_obs is None, only plots real data.
"""
# Joint names
joint_names = [
'L_hip_yaw', 'L_hip_roll', 'L_hip_pitch', 'L_knee', 'L_ankle',
'neck_pitch', 'head_pitch', 'head_yaw', 'head_roll',
'R_hip_yaw', 'R_hip_roll', 'R_hip_pitch', 'R_knee', 'R_ankle'
]
obs_dim = real_obs.shape[1] if sim_obs is None else min(real_obs.shape[1], sim_obs.shape[1])
# Velocity (51D): ang_vel (3) + proj_grav (3) + joint_pos (14) + joint_vel (14) + actions (14) + command (3)
base_ang_vel_start = 0
gravity_start = 3
joint_pos_start = 6
joint_vel_start = 20
action_start = 34
# Create subplot titles with sections
subplot_titles = []
# Base angular velocity (3)
subplot_titles.extend(['<b>BASE ANG VEL</b><br>ω_x', 'ω_y', 'ω_z', ''])
# Raw accelero (3)
subplot_titles.extend(['<b>Raw Accelero</b><br>g_x', 'g_y', 'g_z', ''])
# Joint positions (14 + 2 empty)
subplot_titles.append(f'<b>JOINT POSITIONS</b><br>{joint_names[0]}')
subplot_titles.extend(joint_names[1:14])
subplot_titles.extend(['', ''])
# Joint velocities (14 + 2 empty)
subplot_titles.append(f'<b>JOINT VELOCITIES</b><br>{joint_names[0]}')
subplot_titles.extend(joint_names[1:14])
subplot_titles.extend(['', ''])
# Actions (14 + 2 empty)
subplot_titles.append(f'<b>ACTIONS</b><br>{joint_names[0]}')
subplot_titles.extend(joint_names[1:14])
subplot_titles.extend(['', ''])
num_rows = 14
fig = make_subplots(
rows=num_rows, cols=4,
subplot_titles=subplot_titles,
vertical_spacing=0.02,
horizontal_spacing=0.05,
row_heights=[1]*num_rows,
)
plot_idx = 0
# Track data for common scaling
command_data = []
def add_traces(row, col, real_data, sim_data=None, y_range=None):
"""Helper to add real and sim traces to a subplot."""
fig.add_trace(
go.Scatter(x=real_ts, y=real_data, name='Real',
line=dict(color='blue', width=1.5),
showlegend=(plot_idx == 0)),
row=row, col=col
)
if sim_data is not None:
fig.add_trace(
go.Scatter(x=sim_ts, y=sim_data, name='Sim',
line=dict(color='red', width=1.5, dash='dash'),
showlegend=(plot_idx == 0)),
row=row, col=col
)
if y_range:
fig.update_yaxes(range=y_range, row=row, col=col)
base_ang_vel_data = []
gravity_data = []
joint_pos_data = []
joint_vel_data = []
action_data = []
# 1. Base angular velocity (3 subplots)
for i in range(3):
row, col = divmod(plot_idx, 4)
row += 1
col += 1
base_ang_vel_data.append(real_obs[:, base_ang_vel_start+i])
if sim_obs is not None:
base_ang_vel_data.append(sim_obs[:, base_ang_vel_start+i])
add_traces(row, col, real_obs[:, base_ang_vel_start+i], None if sim_obs is None else sim_obs[:, base_ang_vel_start+i])
fig.update_yaxes(title_text='rad/s', row=row, col=col)
plot_idx += 1
# Empty slot
plot_idx += 1
# 2. Raw accelero (3 subplots)
for i in range(3):
row, col = divmod(plot_idx, 4)
row += 1
col += 1
gravity_data.append(real_obs[:, gravity_start+i])
if sim_obs is not None:
gravity_data.append(sim_obs[:, gravity_start+i])
add_traces(row, col, real_obs[:, gravity_start+i], None if sim_obs is None else sim_obs[:, gravity_start+i])
fig.update_yaxes(title_text='g', row=row, col=col)
plot_idx += 1
# Empty slot
plot_idx += 1
# 3. Joint positions (14 subplots)
for i in range(14):
row, col = divmod(plot_idx, 4)
row += 1
col += 1
if joint_pos_start + i < obs_dim:
joint_pos_data.append(real_obs[:, joint_pos_start+i])
if sim_obs is not None:
joint_pos_data.append(sim_obs[:, joint_pos_start+i])
add_traces(row, col, real_obs[:, joint_pos_start+i], None if sim_obs is None else sim_obs[:, joint_pos_start+i])
fig.update_yaxes(title_text='rad', row=row, col=col)
plot_idx += 1
# Skip 2 empty slots
plot_idx += 2
# 4. Joint velocities (14 subplots)
for i in range(14):
row, col = divmod(plot_idx, 4)
row += 1
col += 1
if joint_vel_start + i < obs_dim:
joint_vel_data.append(real_obs[:, joint_vel_start+i])
if sim_obs is not None:
joint_vel_data.append(sim_obs[:, joint_vel_start+i])
add_traces(row, col, real_obs[:, joint_vel_start+i], None if sim_obs is None else sim_obs[:, joint_vel_start+i])
fig.update_yaxes(title_text='rad/s', row=row, col=col)
plot_idx += 1
# Skip 2 empty slots
plot_idx += 2
# 5. Actions (14 subplots)
for i in range(14):
row, col = divmod(plot_idx, 4)
row += 1
col += 1
if action_start + i < obs_dim:
action_data.append(real_obs[:, action_start+i])
if sim_obs is not None:
action_data.append(sim_obs[:, action_start+i])
add_traces(row, col, real_obs[:, action_start+i], None if sim_obs is None else sim_obs[:, action_start+i])
fig.update_yaxes(title_text='action', row=row, col=col)
fig.update_xaxes(title_text='Time (s)', row=row, col=col)
plot_idx += 1
# Set common y-ranges for each group
def compute_range(data_list):
if not data_list:
return None
all_data = np.concatenate([d.flatten() for d in data_list])
y_min, y_max = np.min(all_data), np.max(all_data)
margin = (y_max - y_min) * 0.1
return [y_min - margin, y_max + margin]
base_ang_vel_range = compute_range(base_ang_vel_data)
gravity_range = compute_range(gravity_data)
joint_pos_range = compute_range(joint_pos_data)
joint_vel_range = compute_range(joint_vel_data)
action_range = compute_range(action_data)
# Apply common ranges
plot_idx = 0
for i in range(3): # Base ang vel
row, col = divmod(plot_idx, 4)
fig.update_yaxes(range=base_ang_vel_range, row=row+1, col=col+1)
plot_idx += 1
plot_idx += 1
for i in range(3): # Gravity
row, col = divmod(plot_idx, 4)
fig.update_yaxes(range=gravity_range, row=row+1, col=col+1)
plot_idx += 1
plot_idx += 1
for i in range(14): # Joint pos
row, col = divmod(plot_idx, 4)
fig.update_yaxes(range=joint_pos_range, row=row+1, col=col+1)
plot_idx += 1
plot_idx += 2
for i in range(14): # Joint vel
row, col = divmod(plot_idx, 4)
fig.update_yaxes(range=joint_vel_range, row=row+1, col=col+1)
plot_idx += 1
plot_idx += 2
for i in range(14): # Actions
row, col = divmod(plot_idx, 4)
fig.update_yaxes(range=action_range, row=row+1, col=col+1)
plot_idx += 1
# Update layout
title = 'Real vs Simulated Observations Comparison' if sim_obs is not None else 'Real Robot Observations'
fig.update_layout(
title_text=title,
title_font_size=24,
height=4600,
width=1600,
showlegend=True,
legend=dict(x=0.85, y=0.99, bgcolor='rgba(255,255,255,0.8)'),
hovermode='x unified'
)
fig.show()
def main():
parser = argparse.ArgumentParser(
description="Compare real and simulated observations (Plotly version)"
)
parser.add_argument("real_pkl", type=str,
help="Path to .pkl file with real robot observations")
parser.add_argument("sim_pkl", type=str, nargs='?', default=None,
help="Path to .pkl file with simulated observations (optional)")
args = parser.parse_args()
# Check if files exist
if not Path(args.real_pkl).exists():
print(f"Error: {args.real_pkl} not found")
return 1
# Load observations
print(f"Loading real observations from {args.real_pkl}...")
real_obs, real_ts = load_observations(args.real_pkl)
print(f"Loaded {len(real_obs)} real observations (shape: {real_obs.shape})")
if args.sim_pkl:
if not Path(args.sim_pkl).exists():
print(f"Error: {args.sim_pkl} not found")
return 1
print(f"Loading simulated observations from {args.sim_pkl}...")
sim_obs, sim_ts = load_observations(args.sim_pkl)
print(f"Loaded {len(sim_obs)} simulated observations (shape: {sim_obs.shape})")
else:
print("No sim data provided, plotting real data only")
sim_obs, sim_ts = None, None
# Plot comparison
print(f"\nGenerating interactive comparison plots...")
plot_comparison(real_obs, real_ts, sim_obs, sim_ts)
return 0
if __name__ == "__main__":
exit(main())
scripts/testbench_sim2real.py
+611
View File
@@ -0,0 +1,611 @@
#!/usr/bin/env python3
"""Sim2real validation for the XL330 test bench.
Runs the same ONNX policy on a fixed sequence of target angles, either in
MuJoCo (with the BAM M6 actuator model) or on the real XL330 (via rustypot),
logs the joint trajectory, and plots sim vs real for comparison.
Example workflow
----------------
# 1) Record in sim:
uv run python scripts/testbench_sim2real.py --mode sim --onnx policy.onnx --out sim.npz
# 2) Plug the real bench via USB, then record on hardware:
uv run python scripts/testbench_sim2real.py --mode real --onnx policy.onnx --out real.npz \
--port /dev/ttyUSB0 --motor-id 1
# 3) Compare the two traces:
uv run python scripts/testbench_sim2real.py --compare sim.npz real.npz --out-plot comparison.png
Observation layout (must match the training env): [joint_pos, joint_vel, last_action, command]
Action: 1-D position offset in radians, scaled by 1.0, added to default pose (0.0).
"""
from __future__ import annotations
import argparse
import math
import time
from pathlib import Path
import numpy as np
import onnxruntime as ort
from mjlab_microduck.robot.testbench_constants import (
TESTBENCH_ARM_MASS,
TESTBENCH_XML,
_set_arm_mass,
)
# --- Match training env ---
CONTROL_DT = 0.02 # decimation=4 × timestep=0.005 (policy rate = 50 Hz)
SIM_DT = 0.005 # (logging rate = 200 Hz — one sample per inner sim step)
LOG_DT = SIM_DT
DEFAULT_POS = 0.0
MAX_ANGLE = math.radians(80.0)
# XL330 present_velocity is returned by rustypot as raw ticks (i32, NOT converted).
# Each tick = 0.229 RPM (per Dynamixel XL330 spec). rad/s = ticks * 0.229 * 2π/60.
DXL_VEL_TICK_TO_RAD_S = 0.229 * 2.0 * math.pi / 60.0 # ≈ 0.02398 rad/s per tick
# ---------------------------------------------------------------------------
# Shared: target schedule + policy wrapper
# ---------------------------------------------------------------------------
def make_target_schedule(
total_time: float,
hold_time: float = 4.0,
seed: int = 0,
) -> np.ndarray:
"""Return one target angle per control step."""
rng = np.random.default_rng(seed)
n_steps = int(round(total_time / CONTROL_DT))
steps_per_hold = int(round(hold_time / CONTROL_DT))
targets = np.zeros(n_steps, dtype=np.float32)
i = 0
while i < n_steps:
angle = float(rng.uniform(-MAX_ANGLE, MAX_ANGLE))
end = min(i + steps_per_hold, n_steps)
targets[i:end] = angle
i = end
return targets
class PolicyRunner:
def __init__(self, onnx_path: str, action_scale: float = 1.0):
print(f"Loading policy: {onnx_path} (action_scale={action_scale})")
self.session = ort.InferenceSession(onnx_path)
self.in_name = self.session.get_inputs()[0].name
in_shape = self.session.get_inputs()[0].shape
print(f" input {self.in_name} shape={in_shape}")
self.action_scale = action_scale
self.last_action = np.zeros(1, dtype=np.float32)
def reset(self):
self.last_action[:] = 0.0
def step(self, q: float, qd: float, target: float) -> float:
# Matches the testbench env's policy obs layout:
# [joint_pos_rel, joint_vel_rel, last_action, command] (4-d).
obs = np.array(
[q - DEFAULT_POS, qd, self.last_action[0], target],
dtype=np.float32,
)[None, :]
action = self.session.run(None, {self.in_name: obs})[0].reshape(-1)
self.last_action = action.astype(np.float32)
return DEFAULT_POS + float(action[0]) * self.action_scale
# ---------------------------------------------------------------------------
# Sim rollout (mujoco, same XL330 testbench XML as training)
# ---------------------------------------------------------------------------
def rollout_sim_bam(onnx_path: str, total_time: float, seed: int, action_scale: float) -> dict:
"""Sim rollout using bam's MujocoController on a vanilla MuJoCo step loop.
Pros: 200 Hz inner-step logging, no torch/mjwarp. Cons: not the exact
actuator that was trained against (uses bam upstream, not mjlab's M6).
"""
import mujoco # local import so --mode real works without mujoco
from bam.actuators import actuators as bam_actuators
from bam.model import models as bam_models
from bam.mujoco import MujocoController
# Load the fitted XL330 m6 params from the canonical bam bundle (identical to
# the values that used to live in mjlab_microduck.actuator.bam_params).
import json as _json
from bam.model import _resolve_json_path
with open(_resolve_json_path(None, "xl330", "m6")) as _f:
DEFAULT_XL330_M6 = _json.load(_f)
VIN = 7.4
KP_FW = 200.0
ACTUATOR_NAME = "1"
# Build BAM's M6 model + XL330 voltage-controlled actuator. The
# MujocoController below drives the joint via this model on every step,
# writing torque to data.ctrl and updating dof_frictionloss/dof_damping
# so MuJoCo's solver applies BAM's Stribeck+load+quadratic friction.
bam_model = bam_models["m6"]()
bam_model.set_actuator(bam_actuators["xl330"]())
bam_model.actuator.kp = KP_FW
bam_model.actuator.vin = VIN
bam_model.load_parameters_from_dict(DEFAULT_XL330_M6)
kt = bam_model.kt.value
R = bam_model.R.value
spec = mujoco.MjSpec.from_file(str(TESTBENCH_XML))
_set_arm_mass(spec, TESTBENCH_ARM_MASS)
# MujocoController needs a torque-controlled motor; the XL330 entry in the
# XML is a position actuator, so convert it and set the voltage-bounded
# force range. Armature is set on the dof by MujocoController.__init__.
for act in spec.actuators:
act.set_to_motor()
act.forcelimited = False
fl = VIN * kt / R
act.forcerange = (-fl, fl)
act.gear = [1.0, 0, 0, 0, 0, 0]
for joint in spec.joints:
if joint.type == mujoco.mjtJoint.mjJNT_HINGE:
joint.damping = 0.0
joint.frictionloss = 0.0
model = spec.compile()
data = mujoco.MjData(model)
model.opt.timestep = SIM_DT
joint_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "1")
dof_id = int(model.jnt_dofadr[joint_id])
qpos_id = int(model.jnt_qposadr[joint_id])
data.qpos[qpos_id] = 0.0
data.qvel[dof_id] = 0.0
mujoco.mj_forward(model, data)
bam_ctrl = MujocoController(bam_model, ACTUATOR_NAME, model, data)
bam_ctrl.reset(data.qpos)
runner = PolicyRunner(onnx_path, action_scale=action_scale)
policy_targets = make_target_schedule(total_time, seed=seed)
decim = int(round(CONTROL_DT / SIM_DT))
# Logging at SIM_DT (200 Hz): decim samples per policy step.
N_log = len(policy_targets) * decim
rec = {k: np.zeros(N_log, dtype=np.float32)
for k in ("t", "target", "q", "qd", "action", "ctrl")}
t = 0.0
log_i = 0
for policy_i, target in enumerate(policy_targets):
q = float(data.qpos[qpos_id])
qd = float(data.qvel[dof_id])
goal = runner.step(q, qd, float(target))
action_raw = float(runner.last_action[0])
for _ in range(decim):
q = float(data.qpos[qpos_id])
dq = float(data.qvel[dof_id])
# ---- log at 200 Hz ----
rec["t"][log_i] = t
rec["target"][log_i] = target
rec["q"][log_i] = q
rec["qd"][log_i] = dq
rec["action"][log_i] = action_raw
rec["ctrl"][log_i] = goal
log_i += 1
# BAM owns control/torque/friction: set the target, then update()
# writes torque to data.ctrl and pushes friction/damping onto the
# dof so MuJoCo's solver applies them on the next step.
bam_ctrl.set_q_target(ACTUATOR_NAME, goal)
bam_ctrl.update()
mujoco.mj_step(model, data)
t += SIM_DT
return rec
def rollout_sim_mjlab(onnx_path: str, total_time: float, seed: int, action_scale: float) -> dict:
"""Sim rollout via the actual mjlab testbench env (same BAM M6 the policy was trained against).
Boots make_testbench_env_cfg() with num_envs=1, overrides the target_angle
command with our deterministic schedule each policy tick, and steps the env
with the policy action. We replicate ManagerBasedRlEnv.step's inner
decimation loop manually so we can log q/qd at SIM_DT (200 Hz) between
sub-steps, matching the bam backend's logging rate.
"""
import torch
from mjlab.envs import ManagerBasedRlEnv
from mjlab_microduck.tasks.testbench_env_cfg import make_testbench_env_cfg
env_cfg = make_testbench_env_cfg(play=True)
env_cfg.scene.num_envs = 1
# Disable auto-resampling and auto-reset so our deterministic schedule and
# initial pose hold for the entire rollout.
env_cfg.commands["target_angle"].resampling_time_range = (1e6, 1e6)
env_cfg.episode_length_s = max(total_time + 10.0, env_cfg.episode_length_s)
# Drop observation noise so the mjlab path is a fair sim2real reference
# (matches the bam path which doesn't inject noise either).
env_cfg.observations["policy"].enable_corruption = False
device = "cuda:0" if torch.cuda.is_available() else "cpu"
env = ManagerBasedRlEnv(cfg=env_cfg, device=device)
env.reset(seed=seed)
cmd_term = env.command_manager.get_term("target_angle")
robot = env.scene["robot"]
runner = PolicyRunner(onnx_path, action_scale=action_scale)
policy_targets = make_target_schedule(total_time, seed=seed)
decim = env.cfg.decimation
physics_dt = env.physics_dt
N_log = len(policy_targets) * decim
rec = {k: np.zeros(N_log, dtype=np.float32)
for k in ("t", "target", "q", "qd", "action", "ctrl")}
t = 0.0
log_i = 0
for target in policy_targets:
# Inject deterministic target and recompute obs so the policy sees it
# this tick (the env's TargetAngleCommand otherwise samples randomly).
cmd_term._target[0, 0] = float(target)
# update_history=True is critical: the testbench env's joint_vel obs
# has a 1-tick delay, so the history buffer must advance each policy
# tick or the policy sees stale velocity.
obs_buf = env.observation_manager.compute(update_history=True)
policy_obs = obs_buf["policy"][0].detach().cpu().numpy().astype(np.float32)
ort_out = runner.session.run(None, {runner.in_name: policy_obs[None, :]})[0].reshape(-1)
runner.last_action = ort_out.astype(np.float32)
action_raw = float(ort_out[0])
goal = DEFAULT_POS + action_raw * action_scale
# Manually run the decimation loop ManagerBasedRlEnv.step uses, so we
# can sample joint state at the physics rate (200 Hz).
action = torch.as_tensor(ort_out, device=device).reshape(1, -1)
env.action_manager.process_action(action)
for _ in range(decim):
# Log the pre-step state to mirror the bam backend (which records
# q/qd right before each mj_step).
rec["t"][log_i] = t
rec["target"][log_i] = float(target)
rec["q"][log_i] = float(robot.data.joint_pos[0, 0].item())
rec["qd"][log_i] = float(robot.data.joint_vel[0, 0].item())
rec["action"][log_i] = action_raw
rec["ctrl"][log_i] = goal
log_i += 1
env.action_manager.apply_action()
env.scene.write_data_to_sim()
env.sim.step()
env.scene.update(dt=physics_dt)
t += physics_dt
env.close()
return rec
# ---------------------------------------------------------------------------
# Real rollout (rustypot XL330)
# ---------------------------------------------------------------------------
def rollout_real(
onnx_path: str,
total_time: float,
seed: int,
port: str,
motor_id: int,
baudrate: int,
kp: int,
action_scale: float,
) -> dict:
from rustypot import Xl330PyController
ctrl = Xl330PyController(port, baudrate, 0.05)
assert ctrl.ping(motor_id), f"motor id={motor_id} not responding on {port}"
# Match the firmware gain used in sim (BAM kp_fw=200).
ctrl.write_torque_enable(motor_id, False)
ctrl.write_operating_mode(motor_id, 3) # position control
ctrl.write_position_p_gain(motor_id, kp)
ctrl.write_position_i_gain(motor_id, 0)
ctrl.write_position_d_gain(motor_id, 0)
# Read back to confirm the gain actually landed (firmware silently clamps
# out-of-range values, so verifying catches mismatches early).
readback = ctrl.read_position_p_gain(motor_id)
if isinstance(readback, (list, tuple)):
readback = readback[0]
print(f" XL330 position P-gain: requested={kp}, readback={readback}")
ctrl.write_goal_position(motor_id, 0.0)
ctrl.write_torque_enable(motor_id, True)
time.sleep(1.0) # let it settle at zero
runner = PolicyRunner(onnx_path, action_scale=action_scale)
policy_targets = make_target_schedule(total_time, seed=seed)
decim = int(round(CONTROL_DT / LOG_DT)) # samples per policy tick (4 at 200 Hz / 50 Hz)
N_log = len(policy_targets) * decim
rec = {k: np.zeros(N_log, dtype=np.float32)
for k in ("t", "target", "q", "qd", "action", "ctrl")}
def _scalar(x) -> float:
if isinstance(x, (list, tuple)):
return float(x[0])
return float(x)
t_start = time.perf_counter()
prev_q = 0.0
log_i = 0
goal = 0.0
action_raw = 0.0
for policy_i, target in enumerate(policy_targets):
tick_start = time.perf_counter()
target_f = float(target)
# Read once, run policy, write goal — all at the start of the 20 ms window.
q = _scalar(ctrl.read_present_position(motor_id))
try:
qd = _scalar(ctrl.read_present_velocity(motor_id)) * DXL_VEL_TICK_TO_RAD_S
except Exception:
qd = (q - prev_q) / CONTROL_DT
goal = runner.step(q, qd, target_f)
action_raw = float(runner.last_action[0])
# ctrl.write_goal_position(motor_id, float(np.clip(goal, -MAX_ANGLE, MAX_ANGLE)))
ctrl.write_goal_position(motor_id, float(goal))
# First 200 Hz sample uses the values we just read (no extra USB round-trip).
rec["t"][log_i] = time.perf_counter() - t_start
rec["target"][log_i] = target_f
rec["q"][log_i] = q
rec["qd"][log_i] = qd
rec["action"][log_i] = action_raw
rec["ctrl"][log_i] = goal
prev_q = q
log_i += 1
# Remaining (decim-1) samples inside the policy window: read only.
for k in range(1, decim):
sample_deadline = tick_start + (k + 1) * LOG_DT
while time.perf_counter() < sample_deadline - 0.001:
time.sleep(0.0005)
q = _scalar(ctrl.read_present_position(motor_id))
try:
qd = _scalar(ctrl.read_present_velocity(motor_id)) * DXL_VEL_TICK_TO_RAD_S
except Exception:
qd = (q - prev_q) / LOG_DT
prev_q = q
rec["t"][log_i] = time.perf_counter() - t_start
rec["target"][log_i] = target_f
rec["q"][log_i] = q
rec["qd"][log_i] = qd
rec["action"][log_i] = action_raw
rec["ctrl"][log_i] = goal
log_i += 1
# Live status on every new segment plus a heartbeat.
new_segment = policy_i == 0 or policy_targets[policy_i] != policy_targets[policy_i - 1]
if new_segment or policy_i % 25 == 0:
print(
f"\r t={rec['t'][log_i-1]:6.2f}s target={math.degrees(target_f):+6.1f}° "
f"q={math.degrees(q):+6.1f}° err={math.degrees(q - target_f):+6.1f}° "
f"goal={math.degrees(goal):+6.1f}°",
end="" if not new_segment else "\n",
flush=True,
)
# Hold the remaining time of the policy window if we got here early.
dt_left = CONTROL_DT - (time.perf_counter() - tick_start)
if dt_left > 0:
time.sleep(dt_left)
print()
ctrl.write_torque_enable(motor_id, False)
return rec
# ---------------------------------------------------------------------------
# Plotting / analytics
# ---------------------------------------------------------------------------
def _mae(a: np.ndarray, b: np.ndarray) -> float:
n = min(len(a), len(b))
return float(np.mean(np.abs(a[:n] - b[:n])))
def npz_to_bam_log(npz_path: str, json_path: str, *, mass: float, length: float,
kp: int, vin: float) -> None:
"""Convert a rollout .npz (written by rollout_sim/rollout_real) to a BAM log json.
BAM log format (see ~/Rhoban/bam/bam/logs.py):
top-level: mass, length, kp, vin, motor, trajectory, dt
entries: position, speed, load, input_volts, temp, goal_position, torque_enable, timestamp
Can be fed to `python -m bam.plot --logdir <dir> --actuator xl330`.
"""
import json
d = dict(np.load(npz_path))
t = d["t"]
# Prefer the actual recorded timestamps for dt to handle small jitter;
# fall back to the fixed control period if there are fewer than 2 samples.
dt = float(np.mean(np.diff(t))) if len(t) > 1 else CONTROL_DT
entries = []
for i in range(len(t)):
entries.append({
"position": float(d["q"][i]),
"speed": float(d["qd"][i]),
"load": 0.0,
"input_volts": vin,
"temp": 25.0,
"goal_position": float(d["ctrl"][i]),
"torque_enable": True,
"timestamp": float(t[i]),
})
log = {
"mass": mass,
"length": length,
"kp": kp,
"vin": vin,
"motor": "xl330",
"trajectory": "rl_policy",
"dt": dt,
"entries": entries,
}
out = Path(json_path)
out.parent.mkdir(parents=True, exist_ok=True)
with open(out, "w") as f:
json.dump(log, f)
print(f"Wrote BAM log: {out} ({len(entries)} entries, dt={dt:.4f}s, mass={mass}kg, kp={kp})")
print(f" Replay with: (cd ~/Rhoban/bam && python -m bam.plot --logdir {out.parent} --actuator xl330)")
def compare_and_plot(sim_file: str, real_file: str, out_path: str) -> None:
import matplotlib.pyplot as plt
sim = dict(np.load(sim_file))
real = dict(np.load(real_file))
n = min(len(sim["t"]), len(real["t"]))
t = sim["t"][:n]
err_sim = sim["q"][:n] - sim["target"][:n]
err_real = real["q"][:n] - real["target"][:n]
print("\n=== Analytics ===")
print(f" steps compared : {n}")
print(f" MAE q (sim vs real) : {_mae(sim['q'], real['q']):.4f} rad "
f"({math.degrees(_mae(sim['q'], real['q'])):.2f}°)")
print(f" sim tracking MAE : {float(np.mean(np.abs(err_sim))):.4f} rad "
f"({math.degrees(float(np.mean(np.abs(err_sim)))):.2f}°)")
print(f" real tracking MAE : {float(np.mean(np.abs(err_real))):.4f} rad "
f"({math.degrees(float(np.mean(np.abs(err_real)))):.2f}°)")
print(f" sim qd RMS : {float(np.sqrt(np.mean(sim['qd'][:n]**2))):.3f} rad/s")
print(f" real qd RMS : {float(np.sqrt(np.mean(real['qd'][:n]**2))):.3f} rad/s")
print(f" action MAE : {_mae(sim['action'], real['action']):.4f} rad")
fig, axes = plt.subplots(4, 1, figsize=(11, 10), sharex=True)
axes[0].plot(t, sim["target"][:n], "k-", lw=1, label="target", alpha=0.4)
axes[0].plot(t, sim["q"][:n], "b-", lw=1.2, label="sim q")
axes[0].plot(t, real["q"][:n], "r-", lw=1.2, label="real q")
axes[0].plot(t, sim["ctrl"][:n], "b:", lw=0.8, alpha=0.6, label="sim goal (policy)")
axes[0].plot(t, real["ctrl"][:n], "r:", lw=0.8, alpha=0.6, label="real goal (policy)")
axes[0].set_ylabel("position [rad]")
axes[0].legend(loc="upper right", fontsize=8)
axes[0].grid(alpha=0.3)
axes[0].set_title(f"Testbench sim2real — MAE(sim, real) = {_mae(sim['q'], real['q']):.4f} rad")
axes[1].plot(t, np.degrees(err_sim), "b-", lw=1, label="sim")
axes[1].plot(t, np.degrees(err_real), "r-", lw=1, label="real")
axes[1].axhline(0, color="k", lw=0.5)
axes[1].set_ylabel("tracking error [deg]")
axes[1].legend(fontsize=8)
axes[1].grid(alpha=0.3)
axes[2].plot(t, sim["qd"][:n], "b-", lw=1, label="sim")
axes[2].plot(t, real["qd"][:n], "r-", lw=1, label="real")
axes[2].set_ylabel("velocity [rad/s]")
axes[2].legend(fontsize=8)
axes[2].grid(alpha=0.3)
axes[3].plot(t, sim["action"][:n], "b-", lw=1, label="sim action")
axes[3].plot(t, real["action"][:n], "r-", lw=1, label="real action")
axes[3].set_ylabel("policy action [rad]")
axes[3].set_xlabel("time [s]")
axes[3].legend(fontsize=8)
axes[3].grid(alpha=0.3)
plt.tight_layout()
plt.savefig(out_path, dpi=140)
print(f"\nSaved plot: {out_path}")
plt.show()
# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--mode", choices=["sim", "real"], help="Rollout mode")
ap.add_argument("--sim-backend", choices=["bam", "mjlab"], default="bam",
help="Sim backend: 'bam' uses bam.MujocoController on a vanilla "
"mujoco loop (200 Hz log, lightweight); 'mjlab' boots the actual "
"make_testbench_env_cfg() mjlab env with its BamM6Actuator (50 Hz log).")
ap.add_argument("--onnx", type=str, help="Path to trained ONNX policy")
ap.add_argument("--out", type=str, help="Output .npz log file")
ap.add_argument("--duration", type=float, default=30.0, help="Total time [s]")
ap.add_argument("--seed", type=int, default=0, help="Target schedule seed")
# real-only
ap.add_argument("--port", type=str, default="/dev/ttyUSB0")
ap.add_argument("--motor-id", type=int, default=1)
ap.add_argument("--baudrate", type=int, default=1_000_000)
ap.add_argument("--kp", type=int, default=200, help="XL330 position P gain")
ap.add_argument("--action-scale", type=float, default=1.0,
help="Multiplier applied to the policy action before offsetting "
"by the default pose (must match training env action scale)")
# compare mode
ap.add_argument("--compare", nargs=2, metavar=("SIM_NPZ", "REAL_NPZ"),
help="Plot two logged runs side by side")
ap.add_argument("--out-plot", type=str, default="testbench_sim2real.png")
# BAM log export
ap.add_argument("--to-bam", nargs=2, metavar=("NPZ", "JSON"),
help="Convert a rollout NPZ to BAM log format "
"(run with: python -m bam.plot --logdir <dir> --actuator xl330)")
ap.add_argument("--bam-mass", type=float, default=TESTBENCH_ARM_MASS, help="Payload mass [kg]")
ap.add_argument("--bam-length", type=float, default=0.1, help="Arm length [m]")
ap.add_argument("--bam-vin", type=float, default=7.4, help="Supply voltage [V]")
args = ap.parse_args()
if args.compare:
compare_and_plot(args.compare[0], args.compare[1], args.out_plot)
return
if args.to_bam:
npz_to_bam_log(
args.to_bam[0],
args.to_bam[1],
mass=args.bam_mass,
length=args.bam_length,
kp=args.kp,
vin=args.bam_vin,
)
return
if not (args.mode and args.onnx and args.out):
ap.error("--mode, --onnx and --out are required for a rollout")
if args.mode == "sim":
sim_fn = rollout_sim_mjlab if args.sim_backend == "mjlab" else rollout_sim_bam
rec = sim_fn(args.onnx, args.duration, args.seed, args.action_scale)
else:
rec = rollout_real(
args.onnx, args.duration, args.seed,
args.port, args.motor_id, args.baudrate, args.kp,
args.action_scale,
)
out = Path(args.out)
out.parent.mkdir(parents=True, exist_ok=True)
np.savez(out, **rec)
err = rec["q"] - rec["target"]
print(f"\nSaved {len(rec['t'])} samples to {out}")
print(f" tracking MAE: {float(np.mean(np.abs(err))):.4f} rad ({math.degrees(float(np.mean(np.abs(err)))):.2f}°)")
if __name__ == "__main__":
main()
scripts/validate_bam_testbench.py
+277
View File
@@ -0,0 +1,277 @@
"""Validate the BAM M6 actuator kernel against real testbench data.
Loads real testbench recordings, replays them in MuJoCo with the BAM M6 actuator,
and compares simulated vs real position traces. Also runs BAM's own Python simulator
as a reference.
Usage:
uv run python3 scripts/validate_bam_testbench.py [--plot] [--max-files N]
"""
import argparse
import json
import os
import sys
from copy import copy
from pathlib import Path
import mujoco
import numpy as np
# ── Paths ──
BAM_DIR = Path(os.path.expanduser("~/Rhoban/bam"))
DATA_DIR = BAM_DIR / "bam" / "data" / "processed"
PARAMS_FILE = BAM_DIR / "params" / "xl330" / "m6_new.json"
TESTBENCH_XML = (
Path(__file__).resolve().parent.parent
/ "src"
/ "mjlab_microduck"
/ "robot"
/ "xl330_test_bench"
/ "scene.xml"
)
# ── Load M6 params ──
with open(PARAMS_FILE) as f:
M6 = json.load(f)
# XL330 firmware constants
ERROR_GAIN = (4096 / (2 * np.pi)) / (256 * 885)
VIN = 7.4
MAX_PWM = 1.0
def bam_python_rollout(log: dict) -> list[float]:
"""Reference: BAM's own Python simulator."""
sys.path.insert(0, str(BAM_DIR))
from bam.model import load_model
from bam.simulate import Simulator
# BAM expects arm_mass in the log (mass of the arm itself, not the payload)
if "arm_mass" not in log:
log = dict(log)
log["arm_mass"] = 0.0
model = load_model(str(PARAMS_FILE))
sim = Simulator(model)
result = sim.rollout_log(log, simulate_control=True)
return result[0] # positions
def compute_m6_friction(motor_torque, external_torque, dq):
"""M6 friction computation matching our kernel (and BAM's model.py)."""
p = M6
stribeck_coeff = np.exp(-(np.abs(dq / p["dtheta_stribeck"]) ** p["alpha"]))
gearbox_torque = np.abs(
external_torque * p["load_friction_external"]
- motor_torque * p["load_friction_motor"]
)
gearbox_torque_stribeck = np.abs(
external_torque * p["load_friction_external_stribeck"]
- motor_torque * p["load_friction_motor_stribeck"]
)
frictionloss = p["friction_base"]
frictionloss += gearbox_torque
frictionloss += stribeck_coeff * p["friction_stribeck"]
frictionloss += gearbox_torque_stribeck * stribeck_coeff
# quadratic (tiny, skip for clarity)
damping = p["friction_viscous"]
friction_budget = frictionloss + damping * np.abs(dq)
return friction_budget
def mujoco_rollout(log: dict) -> list[float]:
"""Run the testbench in MuJoCo with our BAM M6 actuator logic."""
mass = log["mass"]
kp_fw = log["kp"]
dt = log["dt"]
entries = log["entries"]
# Load and modify the testbench model
spec = mujoco.MjSpec.from_file(str(TESTBENCH_XML))
# Convert actuator to motor (same as our kernel's edit_spec)
for act in spec.actuators:
act.set_to_motor()
act.forcelimited = True
force_limit = VIN * M6["kt"] / M6["R"]
act.forcerange = (-force_limit, force_limit)
act.gear = [1.0, 0, 0, 0, 0, 0]
# Zero out MuJoCo joint friction (we handle it)
for joint in spec.joints:
if joint.type == mujoco.mjtJoint.mjJNT_HINGE:
joint.damping = 0.0
joint.frictionloss = 0.0
joint.armature = M6["armature"]
# Set the arm mass to match the BAM recording
for body in spec.bodies:
if body.name == "arm":
# Scale mass and inertia proportionally
original_mass = body.mass
scale = mass / original_mass if original_mass > 0 else 1.0
body.mass = mass
# Scale inertia proportionally to mass
body.fullinertia = [x * scale for x in body.fullinertia]
break
model = spec.compile()
data = mujoco.MjData(model)
model.opt.timestep = dt
# Find joint and actuator IDs
joint_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "1")
dof_id = model.jnt_dofadr[joint_id]
# Initialize state
data.qpos[dof_id] = entries[0]["position"]
data.qvel[dof_id] = entries[0].get("speed", 0.0)
mujoco.mj_forward(model, data)
positions = []
for entry in entries:
positions.append(float(data.qpos[dof_id]))
if not entry["torque_enable"]:
data.ctrl[0] = 0.0
mujoco.mj_step(model, data)
continue
goal = entry["goal_position"]
q = data.qpos[dof_id]
dq = data.qvel[dof_id]
# ── BAM M6 actuator logic (same as our kernel) ──
# 1. Firmware control law
duty = (goal - q) * kp_fw * ERROR_GAIN
duty = np.clip(duty, -MAX_PWM, MAX_PWM)
voltage = VIN * duty
# 2. DC motor torque
motor_torque = M6["kt"] * voltage / M6["R"] - M6["kt"] ** 2 * dq / M6["R"]
# 3. External torque (from MuJoCo bias forces)
# BAM convention: bias_torque = m*g*l*sin(q) with g=-9.81 (gravity negative)
# MuJoCo convention: qfrc_bias has opposite sign
external_torque = -data.qfrc_bias[dof_id]
# 4. M6 friction
friction_budget = compute_m6_friction(motor_torque, external_torque, dq)
# 5. Static friction clipping
eff_inertia = 1.0 / model.dof_invweight0[dof_id] if model.dof_invweight0[dof_id] > 0 else 1e6
net_no_friction = motor_torque + external_torque
tau_stop = (eff_inertia / dt) * dq + net_no_friction
friction_mag = min(abs(tau_stop), friction_budget)
friction_torque = -np.sign(tau_stop) * friction_mag
# 6. Set ctrl = motor + friction (MuJoCo adds qfrc_bias)
data.ctrl[0] = motor_torque + friction_torque
mujoco.mj_step(model, data)
return positions
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--plot", action="store_true", help="Show plots")
parser.add_argument("--max-files", type=int, default=5)
args = parser.parse_args()
data_files = sorted(DATA_DIR.glob("*.json"))
if args.max_files:
data_files = data_files[: args.max_files]
print(f"Validating BAM M6 kernel against {len(data_files)} testbench recordings")
print(f"M6 params: kt={M6['kt']:.4f} R={M6['R']:.4f}")
print(f"Testbench XML: {TESTBENCH_XML}")
print()
results = []
for fpath in data_files:
log = json.load(open(fpath))
name = f"{log['trajectory']}_m{log['mass']}_kp{log['kp']}"
print(f" {name}...", end=" ", flush=True)
real_pos = [e["position"] for e in log["entries"]]
# BAM Python reference
bam_pos = bam_python_rollout(log)
# Our MuJoCo M6 kernel
mj_pos = mujoco_rollout(log)
# Compute MAE
real_np = np.array(real_pos)
bam_np = np.array(bam_pos)
mj_np = np.array(mj_pos[: len(real_np)])
mae_bam = np.mean(np.abs(bam_np - real_np))
mae_mj = np.mean(np.abs(mj_np - real_np))
mae_bam_vs_mj = np.mean(np.abs(bam_np - mj_np))
print(
f"MAE bam_vs_real={mae_bam:.5f} mj_vs_real={mae_mj:.5f} bam_vs_mj={mae_bam_vs_mj:.5f}"
)
results.append(
{
"name": name,
"real": real_np,
"bam": bam_np,
"mj": mj_np,
"mae_bam": mae_bam,
"mae_mj": mae_mj,
"mae_bam_vs_mj": mae_bam_vs_mj,
}
)
print()
avg_bam = np.mean([r["mae_bam"] for r in results])
avg_mj = np.mean([r["mae_mj"] for r in results])
avg_diff = np.mean([r["mae_bam_vs_mj"] for r in results])
print(f"Average MAE bam_vs_real={avg_bam:.5f} mj_vs_real={avg_mj:.5f} bam_vs_mj={avg_diff:.5f}")
if avg_diff > 0.01:
print("\n⚠ BAM and MuJoCo diverge significantly — likely a kernel bug!")
elif avg_mj > avg_bam * 1.5:
print("\n⚠ MuJoCo worse than BAM — MuJoCo dynamics differ from BAM's simple integrator")
else:
print("\n✓ BAM and MuJoCo agree — kernel is correct")
if args.plot:
try:
import matplotlib.pyplot as plt
n = len(results)
fig, axes = plt.subplots(n, 1, figsize=(12, 3 * n), sharex=False)
if n == 1:
axes = [axes]
for ax, r in zip(axes, results):
t = np.arange(len(r["real"])) * 0.005
ax.plot(t, r["real"], "k-", lw=1.5, label="Real")
ax.plot(t, r["bam"], "b--", lw=1.2, label=f'BAM (MAE={r["mae_bam"]:.4f})')
ax.plot(t, r["mj"], "r:", lw=1.2, label=f'MuJoCo M6 (MAE={r["mae_mj"]:.4f})')
ax.set_title(r["name"])
ax.set_ylabel("Position (rad)")
ax.legend(fontsize=8)
ax.grid(alpha=0.3)
axes[-1].set_xlabel("Time (s)")
plt.tight_layout()
plt.savefig("bam_validation.png", dpi=150)
print("Saved bam_validation.png")
plt.show()
except ImportError:
print("matplotlib not available, skipping plots")
if __name__ == "__main__":
main()
scripts/view_slope_terrain.py
+101
View File
@@ -0,0 +1,101 @@
"""Observer la rampe du mode pente (roller_slope) dans le viewer MuJoCo.
Construit UNIQUEMENT le terrain « plat + rampe » (FlatRampTerrainCfg), sur
plusieurs rangées de difficulté croissante (raideur 2° -> 20°), et ouvre le
viewer natif MuJoCo. Aucune politique entraînée n'est nécessaire — c'est fait
pour valider à l'œil la géométrie (jointure plat/rampe, sens de descente).
Usage :
uv run python scripts/view_slope_terrain.py
uv run python scripts/view_slope_terrain.py --rows 6 --ramp-max 8 --runout 4
uv run python scripts/view_slope_terrain.py --build-only # test sans GUI
Dans le viewer : molette pour zoomer, clic-gauche glisser pour orbiter,
clic-droit glisser pour translater. Chaque rangée est une rampe de plus en plus
raide (difficulté 0 -> 1), de longueur tirée au hasard dans [ramp-min, ramp-max],
et se termine par un plat de sortie. « Devant » (+x) doit descendre.
"""
import argparse
import mujoco
import mujoco.viewer
from mjlab.terrains.terrain_generator import TerrainGenerator, TerrainGeneratorCfg
from mjlab_microduck.tasks.slope_terrain import (
FlatRampTerrainCfg,
RAMP_DEG_MIN,
RAMP_DEG_MAX,
)
def build_model(rows, size, flat_length, ramp_range, runout, deg_min, deg_max):
"""Construit le modèle MuJoCo du terrain seul (rows rampes de raideur croissante)."""
cfg = TerrainGeneratorCfg(
seed=0,
size=size,
num_rows=rows,
num_cols=1,
curriculum=True, # difficulté croissante le long des rangées
difficulty_range=(0.0, 1.0),
add_lights=True,
sub_terrains={
"flat_ramp": FlatRampTerrainCfg(
flat_length=flat_length,
ramp_length_range=ramp_range,
runout_length=runout,
deg_min=deg_min,
deg_max=deg_max,
)
},
)
generator = TerrainGenerator(cfg)
spec = mujoco.MjSpec()
generator.compile(spec)
model = spec.compile()
return model
def main():
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
p.add_argument("--rows", type=int, default=5, help="Nb de rangées = nb de raideurs affichées (défaut 5)")
p.add_argument("--size", type=float, nargs=2, default=(15.0, 4.0), help="Taille d'une tuile (x y) en m")
p.add_argument("--flat-length", type=float, default=2.0, help="Longueur du plat de départ (m)")
p.add_argument("--ramp-min", type=float, default=3.0, help="Longueur horizontale mini de la rampe (m)")
p.add_argument("--ramp-max", type=float, default=8.0, help="Longueur horizontale maxi de la rampe (m)")
p.add_argument("--runout", type=float, default=4.0, help="Longueur du plat de sortie (m)")
p.add_argument("--deg-min", type=float, default=RAMP_DEG_MIN, help=f"Raideur min en degrés (défaut {RAMP_DEG_MIN})")
p.add_argument("--deg-max", type=float, default=RAMP_DEG_MAX, help=f"Raideur max en degrés (défaut {RAMP_DEG_MAX})")
p.add_argument("--build-only", action="store_true", help="Construit le modèle et quitte (test sans GUI)")
args = p.parse_args()
model = build_model(
rows=args.rows,
size=tuple(args.size),
flat_length=args.flat_length,
ramp_range=(args.ramp_min, args.ramp_max),
runout=args.runout,
deg_min=args.deg_min,
deg_max=args.deg_max,
)
data = mujoco.MjData(model)
mujoco.mj_forward(model, data)
print(
f"Terrain construit : {args.rows} rampes, raideur {args.deg_min}°->{args.deg_max}°, "
f"longueur rampe {args.ramp_min}-{args.ramp_max}m + sortie {args.runout}m, "
f"{model.ngeom} géométries."
)
if args.build_only:
print("--build-only : OK, pas de GUI.")
return
print("Ouverture du viewer MuJoCo (Ctrl+C pour quitter)…")
with mujoco.viewer.launch_passive(model, data, show_left_ui=False, show_right_ui=False) as viewer:
while viewer.is_running():
mujoco.mj_forward(model, data)
viewer.sync()
if __name__ == "__main__":
main()
scripts/wandb_utils.py
+138
View File
@@ -0,0 +1,138 @@
"""Shared wandb helpers for play_latest and export_latest scripts."""
import os
import subprocess
import sys
from datetime import datetime
import wandb
WANDB_PROJECT = "pollen-robotics/mjlab_microduck"
def find_latest_run(user_filter: str) -> wandb.apis.public.Run | None:
"""Return the most recent run whose metadata email contains *user_filter*."""
api = wandb.Api()
runs = api.runs(WANDB_PROJECT, per_page=100, order="-created_at")
for run in runs:
email = (run.metadata or {}).get("email", "")
if user_filter.lower() in email.lower():
return run
return None
def find_latest_runs(user_filter, task_match, n):
"""Return up to *n* most recent runs whose metadata email contains
*user_filter* and whose task_id (metadata args[0]) satisfies *task_match*.
task_match: Callable[[str], bool] applied to the wandb task_id string.
Ordered most-recent first.
"""
api = wandb.Api()
runs = api.runs(WANDB_PROJECT, per_page=100, order="-created_at")
matched = []
for run in runs:
email = (run.metadata or {}).get("email", "")
if user_filter.lower() not in email.lower():
continue
args = (run.metadata or {}).get("args", [])
task_id = args[0] if args else ""
if not task_match(task_id):
continue
matched.append(run)
if len(matched) >= n:
break
return matched
def _format_duration(seconds: float) -> str:
h = int(seconds // 3600)
m = int((seconds % 3600) // 60)
return f"{h}h{m:02d}m"
def print_run_info(run: wandb.apis.public.Run) -> dict:
"""Print run details and return dict with env_name, run_path, checkpoints."""
meta = run.metadata
summary = run.summary
train_cfg = run.config.get("train_cfg", {})
env_cfg = run.config.get("env_cfg", {})
env_name = meta.get("args", ["?"])[0]
run_path = f"{WANDB_PROJECT}/{run.id}"
duration = summary.get("_runtime", 0)
total_steps = summary.get("_step", 0)
max_iter = train_cfg.get("max_iterations", "?")
num_envs = env_cfg.get("scene", {}).get("terrain", {}).get("num_envs", "?")
mean_reward = summary.get("Train/mean_reward")
lr = summary.get("Loss/learning_rate")
reward_items = [
(k.removeprefix("Episode_Reward/"), v)
for k, v in summary.items()
if k.startswith("Episode_Reward/") and isinstance(v, (int, float))
]
reward_items.sort(key=lambda x: abs(x[1]), reverse=True)
checkpoints = sorted(
[f.name for f in run.files() if f.name.startswith("model_") and f.name.endswith(".pt")]
)
last_ckpt = checkpoints[-1] if checkpoints else "none"
created = datetime.fromisoformat(run.created_at.replace("Z", "+00:00"))
date_str = created.astimezone().strftime("%Y-%m-%d %H:%M")
print("=" * 60)
print(f" Run: {run.name}")
print(f" ID: {run.id}")
print(f" Date: {date_str}")
print(f" Status: {run.state}")
print(f" Environment: {env_name}")
print(f" Duration: {_format_duration(duration)}")
print(f" Progress: {total_steps} / {max_iter} iterations")
print(f" Num envs: {num_envs}")
print(f" Mean reward: {mean_reward:.2f}" if mean_reward else " Mean reward: ?")
print(f" Learning rate:{lr:.2e}" if lr else " Learning rate:?")
print(f" Last ckpt: {last_ckpt}")
print(f" Host: {meta.get('host', '?')}")
print(f" GPU: {meta.get('gpu', '?')}")
print(f" User: {meta.get('email', '?')}")
print()
print(" Top rewards (by magnitude):")
for name, val in reward_items[:8]:
sign = "+" if val >= 0 else ""
print(f" {name:<35s} {sign}{val:.4f}")
print("=" * 60)
return {"env_name": env_name, "run_path": run_path, "checkpoints": checkpoints}
def resolve_run(user: str, task_substr: str | None = None) -> tuple[wandb.apis.public.Run, dict]:
"""Find latest run for user (optionally filtered to a task-id substring),
print info, return (run, info). Exits on failure."""
if task_substr:
print(f"Searching latest '{task_substr}' run for user '{user}'...")
runs = find_latest_runs(user, lambda t: task_substr.lower() in t.lower(), 1)
run = runs[0] if runs else None
else:
print(f"Searching latest run for user '{user}'...")
run = find_latest_run(user)
if run is None:
suffix = f" matching '{task_substr}'" if task_substr else ""
print(f"No run found for user '{user}'{suffix}", file=sys.stderr)
sys.exit(1)
info = print_run_info(run)
return run, info
def run_command(cmd: list[str], dry_run: bool) -> None:
"""Print and optionally execute a command from the project root."""
print()
print("Command:")
print(f" {' '.join(cmd)}")
print()
if dry_run:
print("(dry-run, not executing)")
return
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
subprocess.run(cmd, cwd=project_root)