"""NumPy 最小二乘学习小车动力学；显式 --run 才拟合并保存完整记录。""" from argparse import ArgumentParser from pathlib import Path import json import platform import random import time import numpy as np from world_models import ( DT, HERE, MODES, PLAN, control, digest, features, predict, rollout, save_new, step, ) def make_plan(): return { "schema": 1, "data_seeds": [11, 22, 33], "train_trajectories": 256, "dev_trajectories": 64, "test_trajectories": 64, "train_steps": 8, "test_steps": 16, "initial_x_range": [-2.0, 2.0], "initial_velocity_range": [-0.6, 0.6], "dt_seconds": DT, "retention": 0.8, "velocity_gain": 0.4, "actuator_saturation": [-1, 1], "train_commands": [-1, 0, 1], "models": list(MODES), "fit": "float64 numpy.linalg.lstsq rcond=None", "observation": "exact positions; no pixels, no noise, no learned encoder", "test_regimes": { "standard": {"retention": 0.8, "command_scale": 1}, "strong_commands": {"retention": 0.8, "command_scale": 3}, "changed_retention": {"retention": 0.3, "command_scale": 1}, }, "control": { "initial_x": [-1.5, 0.0, 1.5], "initial_velocity": [-0.6, 0.0, 0.6], "goals": [-1.0, 1.0], "steps": 24, "planning_horizon": 4, "commands": [0, -1, 1], "wide_commands": [0, -3, 3], "terminal_velocity_weight": 0.1, "command_cost": 0.01, "tie_round_digits": 12, "position_threshold": 0.15, "speed_threshold": 0.2, }, "selection": "all seeds, all models, no hyperparameter search", } def trajectory(seed, split, index, steps, scale=1, retention=0.8): offsets = {"train": 10000, "dev": 20000, "test": 30000} rng = random.Random(seed * 1000000 + offsets[split] + index) x, velocity = rng.uniform(-2, 2), rng.uniform(-0.6, 0.6) previous = x - DT * velocity initial = {"previous": previous, "x": x, "velocity": velocity} observations, velocities, actions = [previous, x], [velocity], [] for _ in range(steps): action = scale * rng.choice([-1, 0, 1]) x, velocity = step(x, velocity, action, retention) observations.append(x) velocities.append(velocity) actions.append(action) return {"id": f"{seed}-{split}-{index}", "initial": initial, "positions": observations, "velocities": velocities, "actions": actions} def rows(trajectories, mode): inputs, targets = [], [] for trajectory in trajectories: p = trajectory["positions"] for t, a in enumerate(trajectory["actions"]): inputs.append(features(mode, p[t], p[t+1], a)) targets.append(p[t+2]) return np.asarray(inputs, dtype=np.float64), np.asarray(targets, dtype=np.float64) def evaluate(model, trajectories): one_errors, errors, predictions = [], [], [] for tr in trajectories: p, actions = tr["positions"], tr["actions"] one = [predict(model, p[t], p[t+1], a) for t, a in enumerate(actions)] many = rollout(model, p[0], p[1], actions) one_errors.extend((np.array(one) - p[2:]) ** 2) errors.append((np.array(many) - p[2:]) ** 2) predictions.append({"id": tr["id"], "one_step": one, "rollout": many}) return {"one_step_rmse": float(np.sqrt(np.mean(one_errors))), "rollout_rmse": np.sqrt(np.mean(errors, axis=0)).tolist(), "predictions": predictions} def run(plan_path, out_dir): plan = json.loads(plan_path.read_text()) assert plan == make_plan(), "配置与代码不一致；先用修改后的脚本重新冻结。" out_dir.mkdir(parents=True, exist_ok=False) start = time.perf_counter() methods = list(MODES) + ["oracle", "history_action_wide", "oracle_wide"] record = {"complete": False, "plan": plan, "plan_sha256": digest(plan_path), "source_sha256": {p.name: digest(p) for p in [HERE/"world_models.py", Path(__file__)]}, "environment": {"python": platform.python_version(), "numpy": np.__version__, "platform": platform.platform()}, "control_methods": methods, "runs": []} for seed in plan["data_seeds"]: training = [trajectory(seed, "train", i, 8) for i in range(256)] dev = [trajectory(seed, "dev", i, 16) for i in range(64)] tests = {name: [trajectory(seed, "test", i, 16, spec["command_scale"], spec["retention"]) for i in range(64)] for name, spec in plan["test_regimes"].items()} models, development, evaluation = {}, {}, {k: {} for k in tests} for mode in MODES: x, y = rows(training, mode) w, residuals, rank, singular = np.linalg.lstsq(x, y, rcond=None) assert rank == x.shape[1] model = {"mode": mode, "weights": w.tolist(), "rank": int(rank), "singular_values": singular.tolist(), "fit_residual_sum": float(np.sum((x @ w - y)**2)), "reported_residuals": residuals.tolist(), "samples": len(y)} models[mode] = model development[mode] = evaluate(model, dev) for name, trajectories in tests.items(): evaluation[name][mode] = evaluate(model, trajectories) cases = [{"id": i, "x": x, "velocity": v, "goal": g} for i, (x, v, g) in enumerate( (x, v, g) for x in [-1.5, 0.0, 1.5] for v in [-0.6, 0.0, 0.6] for g in [-1.0, 1.0])] controls = {} for name in methods: base = name.removesuffix("_wide") model = {"mode": "oracle"} if base == "oracle" else models[base] actions = (0, -3, 3) if name.endswith("_wide") else (0, -1, 1) controls[name] = [control(model, case, actions) for case in cases] model_path = out_dir / f"seed-{seed}-models.json" save_new(model_path, models) assert json.loads(model_path.read_text()) == models record["runs"].append({ "seed": seed, "training": training, "dev": dev, "tests": tests, "models": models, "model_file": model_path.name, "model_sha256": digest(model_path), "development": development, "evaluation": evaluation, "control": controls, }) print("完成数据种子", seed, flush=True) record["wall_seconds"] = time.perf_counter() - start record["complete"] = True save_new(out_dir/"record.json", record) print(out_dir/"record.json") def main(): parser = ArgumentParser(description=__doc__) parser.add_argument("--freeze", type=Path) parser.add_argument("--run", action="store_true") parser.add_argument("--plan", type=Path, default=PLAN) parser.add_argument("--out-dir", type=Path) args = parser.parse_args() if args.freeze: save_new(args.freeze, make_plan()) print(args.freeze) elif args.run: if args.out_dir is None: parser.error("--run 需要 --out-dir 新目录") run(args.plan, args.out_dir) else: print("默认不拟合。使用 --run --out-dir 新目录，运行三份数据上的九次拟合。") if __name__ == "__main__": main()