♻️ crazy refactor
This commit is contained in:
376
src/sysid/optimize.py
Normal file
376
src/sysid/optimize.py
Normal file
@@ -0,0 +1,376 @@
|
||||
"""CMA-ES optimiser — fit simulation parameters to a real-robot recording.
|
||||
|
||||
Minimises the trajectory-matching cost between a MuJoCo rollout and a
|
||||
recorded real-robot sequence. Uses the ``cmaes`` package (pure-Python
|
||||
CMA-ES with native box-constraint support).
|
||||
|
||||
Usage:
|
||||
python -m src.sysid.optimize \
|
||||
--robot-path assets/rotary_cartpole \
|
||||
--recording assets/rotary_cartpole/recordings/capture_20260311_120000.npz
|
||||
|
||||
# Shorter run for testing:
|
||||
python -m src.sysid.optimize \
|
||||
--robot-path assets/rotary_cartpole \
|
||||
--recording <file>.npz \
|
||||
--max-generations 10 --population-size 8
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import time
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import structlog
|
||||
|
||||
from src.sysid.rollout import (
|
||||
ROTARY_CARTPOLE_PARAMS,
|
||||
ParamSpec,
|
||||
bounds_arrays,
|
||||
defaults_vector,
|
||||
params_to_dict,
|
||||
rollout,
|
||||
windowed_rollout,
|
||||
)
|
||||
|
||||
log = structlog.get_logger()
|
||||
|
||||
|
||||
# ── Cost function ────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _angle_diff(a: np.ndarray, b: np.ndarray) -> np.ndarray:
|
||||
"""Shortest signed angle difference, handling wrapping."""
|
||||
return np.arctan2(np.sin(a - b), np.cos(a - b))
|
||||
|
||||
|
||||
def _check_inertia_valid(params: dict[str, float]) -> bool:
|
||||
"""Quick reject: pendulum inertia tensor must be positive-definite."""
|
||||
ixx = params.get("pendulum_ixx", 6.16e-06)
|
||||
iyy = params.get("pendulum_iyy", 6.16e-06)
|
||||
izz = params.get("pendulum_izz", 1.23e-05)
|
||||
ixy = params.get("pendulum_ixy", 6.10e-06)
|
||||
det_xy = ixx * iyy - ixy * ixy
|
||||
return det_xy > 0 and ixx > 0 and iyy > 0 and izz > 0
|
||||
|
||||
|
||||
def _compute_trajectory_cost(
|
||||
sim: dict[str, np.ndarray],
|
||||
recording: dict[str, np.ndarray],
|
||||
pos_weight: float = 1.0,
|
||||
vel_weight: float = 0.1,
|
||||
) -> float:
|
||||
"""Weighted MSE between sim and real trajectories."""
|
||||
motor_err = _angle_diff(sim["motor_angle"], recording["motor_angle"])
|
||||
pend_err = _angle_diff(sim["pendulum_angle"], recording["pendulum_angle"])
|
||||
motor_vel_err = sim["motor_vel"] - recording["motor_vel"]
|
||||
pend_vel_err = sim["pendulum_vel"] - recording["pendulum_vel"]
|
||||
|
||||
return float(
|
||||
pos_weight * np.mean(motor_err**2)
|
||||
+ pos_weight * np.mean(pend_err**2)
|
||||
+ vel_weight * np.mean(motor_vel_err**2)
|
||||
+ vel_weight * np.mean(pend_vel_err**2)
|
||||
)
|
||||
|
||||
|
||||
def cost_function(
|
||||
params_vec: np.ndarray,
|
||||
recording: dict[str, np.ndarray],
|
||||
robot_path: Path,
|
||||
specs: list[ParamSpec],
|
||||
sim_dt: float = 0.002,
|
||||
substeps: int = 10,
|
||||
pos_weight: float = 1.0,
|
||||
vel_weight: float = 0.1,
|
||||
window_duration: float = 0.5,
|
||||
) -> float:
|
||||
"""Compute trajectory-matching cost for a candidate parameter vector.
|
||||
|
||||
Uses **multiple-shooting** (windowed rollout): the recording is split
|
||||
into short windows (default 0.5 s). Each window is initialised from
|
||||
the real qpos/qvel, so early errors don’t compound across the full
|
||||
trajectory. This gives a much smoother cost landscape for CMA-ES.
|
||||
|
||||
Set ``window_duration=0`` to fall back to the original open-loop
|
||||
single-shot rollout (not recommended).
|
||||
"""
|
||||
params = params_to_dict(params_vec, specs)
|
||||
|
||||
if not _check_inertia_valid(params):
|
||||
return 1e6
|
||||
|
||||
try:
|
||||
if window_duration > 0:
|
||||
sim = windowed_rollout(
|
||||
robot_path=robot_path,
|
||||
params=params,
|
||||
recording=recording,
|
||||
window_duration=window_duration,
|
||||
sim_dt=sim_dt,
|
||||
substeps=substeps,
|
||||
)
|
||||
else:
|
||||
sim = rollout(
|
||||
robot_path=robot_path,
|
||||
params=params,
|
||||
actions=recording["action"],
|
||||
timesteps=recording["time"],
|
||||
sim_dt=sim_dt,
|
||||
substeps=substeps,
|
||||
)
|
||||
except Exception as exc:
|
||||
log.warning("rollout_failed", error=str(exc))
|
||||
return 1e6
|
||||
|
||||
return _compute_trajectory_cost(sim, recording, pos_weight, vel_weight)
|
||||
|
||||
|
||||
# ── CMA-ES optimisation loop ────────────────────────────────────────
|
||||
|
||||
|
||||
def optimize(
|
||||
robot_path: str | Path,
|
||||
recording_path: str | Path,
|
||||
specs: list[ParamSpec] | None = None,
|
||||
sigma0: float = 0.3,
|
||||
population_size: int = 20,
|
||||
max_generations: int = 1000,
|
||||
sim_dt: float = 0.002,
|
||||
substeps: int = 10,
|
||||
pos_weight: float = 1.0,
|
||||
vel_weight: float = 0.1,
|
||||
window_duration: float = 0.5,
|
||||
seed: int = 42,
|
||||
) -> dict:
|
||||
"""Run CMA-ES optimisation and return results.
|
||||
|
||||
Returns a dict with:
|
||||
best_params: dict[str, float]
|
||||
best_cost: float
|
||||
history: list of (generation, best_cost) tuples
|
||||
recording: str (path used)
|
||||
specs: list of param names
|
||||
"""
|
||||
from cmaes import CMA
|
||||
|
||||
robot_path = Path(robot_path).resolve()
|
||||
recording_path = Path(recording_path).resolve()
|
||||
|
||||
if specs is None:
|
||||
specs = ROTARY_CARTPOLE_PARAMS
|
||||
|
||||
# Load recording.
|
||||
recording = dict(np.load(recording_path))
|
||||
n_samples = len(recording["time"])
|
||||
duration = recording["time"][-1] - recording["time"][0]
|
||||
n_windows = max(1, int(duration / window_duration)) if window_duration > 0 else 1
|
||||
log.info(
|
||||
"recording_loaded",
|
||||
path=str(recording_path),
|
||||
samples=n_samples,
|
||||
duration=f"{duration:.1f}s",
|
||||
window_duration=f"{window_duration}s",
|
||||
n_windows=n_windows,
|
||||
)
|
||||
|
||||
# Initial point (defaults) — normalised to [0, 1] for CMA-ES.
|
||||
lo, hi = bounds_arrays(specs)
|
||||
x0 = defaults_vector(specs)
|
||||
|
||||
# Normalise to [0, 1] for the optimizer (better conditioned).
|
||||
span = hi - lo
|
||||
span[span == 0] = 1.0 # avoid division by zero
|
||||
|
||||
def to_normed(x: np.ndarray) -> np.ndarray:
|
||||
return (x - lo) / span
|
||||
|
||||
def from_normed(x_n: np.ndarray) -> np.ndarray:
|
||||
return x_n * span + lo
|
||||
|
||||
x0_normed = to_normed(x0)
|
||||
bounds_normed = np.column_stack(
|
||||
[np.zeros(len(specs)), np.ones(len(specs))]
|
||||
)
|
||||
|
||||
optimizer = CMA(
|
||||
mean=x0_normed,
|
||||
sigma=sigma0,
|
||||
bounds=bounds_normed,
|
||||
population_size=population_size,
|
||||
seed=seed,
|
||||
)
|
||||
|
||||
best_cost = float("inf")
|
||||
best_params_vec = x0.copy()
|
||||
history: list[tuple[int, float]] = []
|
||||
|
||||
log.info(
|
||||
"cmaes_starting",
|
||||
n_params=len(specs),
|
||||
population=population_size,
|
||||
max_gens=max_generations,
|
||||
sigma0=sigma0,
|
||||
)
|
||||
|
||||
t0 = time.monotonic()
|
||||
|
||||
for gen in range(max_generations):
|
||||
solutions = []
|
||||
for _ in range(optimizer.population_size):
|
||||
x_normed = optimizer.ask()
|
||||
x_natural = from_normed(x_normed)
|
||||
|
||||
# Clip to bounds (CMA-ES can slightly exceed with sampling noise).
|
||||
x_natural = np.clip(x_natural, lo, hi)
|
||||
|
||||
c = cost_function(
|
||||
x_natural,
|
||||
recording,
|
||||
robot_path,
|
||||
specs,
|
||||
sim_dt=sim_dt,
|
||||
substeps=substeps,
|
||||
pos_weight=pos_weight,
|
||||
vel_weight=vel_weight,
|
||||
window_duration=window_duration,
|
||||
)
|
||||
solutions.append((x_normed, c))
|
||||
|
||||
if c < best_cost:
|
||||
best_cost = c
|
||||
best_params_vec = x_natural.copy()
|
||||
|
||||
optimizer.tell(solutions)
|
||||
history.append((gen, best_cost))
|
||||
|
||||
elapsed = time.monotonic() - t0
|
||||
if gen % 5 == 0 or gen == max_generations - 1:
|
||||
log.info(
|
||||
"cmaes_generation",
|
||||
gen=gen,
|
||||
best_cost=f"{best_cost:.6f}",
|
||||
elapsed=f"{elapsed:.1f}s",
|
||||
gen_best=f"{min(c for _, c in solutions):.6f}",
|
||||
)
|
||||
|
||||
total_time = time.monotonic() - t0
|
||||
best_params = params_to_dict(best_params_vec, specs)
|
||||
|
||||
log.info(
|
||||
"cmaes_finished",
|
||||
best_cost=f"{best_cost:.6f}",
|
||||
total_time=f"{total_time:.1f}s",
|
||||
evaluations=max_generations * population_size,
|
||||
)
|
||||
|
||||
# Log parameter comparison.
|
||||
defaults = params_to_dict(defaults_vector(specs), specs)
|
||||
for name in best_params:
|
||||
d = defaults[name]
|
||||
b = best_params[name]
|
||||
change_pct = ((b - d) / abs(d) * 100) if abs(d) > 1e-12 else 0.0
|
||||
log.info(
|
||||
"param_result",
|
||||
name=name,
|
||||
default=f"{d:.6g}",
|
||||
tuned=f"{b:.6g}",
|
||||
change=f"{change_pct:+.1f}%",
|
||||
)
|
||||
|
||||
return {
|
||||
"best_params": best_params,
|
||||
"best_cost": best_cost,
|
||||
"history": history,
|
||||
"recording": str(recording_path),
|
||||
"param_names": [s.name for s in specs],
|
||||
"defaults": {s.name: s.default for s in specs},
|
||||
"timestamp": datetime.now().isoformat(),
|
||||
}
|
||||
|
||||
|
||||
# ── CLI entry point ──────────────────────────────────────────────────
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Fit simulation parameters to a real-robot recording (CMA-ES)."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--robot-path",
|
||||
type=str,
|
||||
default="assets/rotary_cartpole",
|
||||
help="Path to robot asset directory",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--recording",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Path to .npz recording file",
|
||||
)
|
||||
parser.add_argument("--sigma0", type=float, default=0.3)
|
||||
parser.add_argument("--population-size", type=int, default=20)
|
||||
parser.add_argument("--max-generations", type=int, default=200)
|
||||
parser.add_argument("--sim-dt", type=float, default=0.002)
|
||||
parser.add_argument("--substeps", type=int, default=10)
|
||||
parser.add_argument("--pos-weight", type=float, default=1.0)
|
||||
parser.add_argument("--vel-weight", type=float, default=0.1)
|
||||
parser.add_argument(
|
||||
"--window-duration",
|
||||
type=float,
|
||||
default=0.5,
|
||||
help="Shooting window length in seconds (0 = open-loop, default 0.5)",
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
parser.add_argument(
|
||||
"--no-export",
|
||||
action="store_true",
|
||||
help="Skip exporting tuned files (results JSON only)",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
result = optimize(
|
||||
robot_path=args.robot_path,
|
||||
recording_path=args.recording,
|
||||
sigma0=args.sigma0,
|
||||
population_size=args.population_size,
|
||||
max_generations=args.max_generations,
|
||||
sim_dt=args.sim_dt,
|
||||
substeps=args.substeps,
|
||||
pos_weight=args.pos_weight,
|
||||
vel_weight=args.vel_weight,
|
||||
window_duration=args.window_duration,
|
||||
seed=args.seed,
|
||||
)
|
||||
|
||||
# Save results JSON.
|
||||
robot_path = Path(args.robot_path).resolve()
|
||||
result_path = robot_path / "sysid_result.json"
|
||||
# Convert numpy types for JSON serialisation.
|
||||
result_json = {
|
||||
k: v for k, v in result.items() if k != "history"
|
||||
}
|
||||
result_json["history_summary"] = {
|
||||
"first_cost": result["history"][0][1] if result["history"] else None,
|
||||
"final_cost": result["history"][-1][1] if result["history"] else None,
|
||||
"generations": len(result["history"]),
|
||||
}
|
||||
result_path.write_text(json.dumps(result_json, indent=2, default=str))
|
||||
log.info("results_saved", path=str(result_path))
|
||||
|
||||
# Export tuned files unless --no-export.
|
||||
if not args.no_export:
|
||||
from src.sysid.export import export_tuned_files
|
||||
|
||||
export_tuned_files(
|
||||
robot_path=args.robot_path,
|
||||
params=result["best_params"],
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user