"""标准库：小车环境、规划器和图像/注意力/跨动作一致性算例。"""
from argparse import ArgumentParser
from hashlib import sha256
from itertools import product
from pathlib import Path
import json

HERE = Path(__file__).resolve().parent
PLAN = HERE / "data/world-models-plan.json"
RECORD = HERE / "data/world-models-fit.json"
SMALL = HERE / "data/world-models-small.json"
MODES = ("position_action", "history_no_action", "history_action")
NAMES = ("单位置＋动作", "两位置，无动作", "两位置＋动作")
DT = 0.2


def digest(path):
    return sha256(Path(path).read_bytes()).hexdigest()


def save_new(path, data):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("x") as f:
        json.dump(data, f, ensure_ascii=False, indent=2, allow_nan=False)
        f.write("\n")


def step(x, velocity, action, retention=0.8):
    """x 米、velocity 米/秒；命令 action 无单位，执行器在 ±1 饱和。"""
    applied = max(-1.0, min(1.0, action))
    next_velocity = retention * velocity + 0.4 * applied
    return x + DT * next_velocity, next_velocity


def features(mode, previous, current, action):
    if mode == "position_action":
        return [current, action, 1.0]
    if mode == "history_no_action":
        return [current, previous, 1.0]
    if mode == "history_action":
        return [current, previous, action, 1.0]
    raise ValueError(mode)


def predict(model, previous, current, action):
    if model["mode"] == "oracle":
        return step(current, (current - previous) / DT, action)[0]
    return sum(w * x for w, x in zip(
        model["weights"], features(model["mode"], previous, current, action),
    ))


def rollout(model, previous, current, actions):
    result = []
    for action in actions:
        following = predict(model, previous, current, action)
        result.append(following)
        previous, current = current, following
    return result


def choose(model, previous, current, goal, actions=(0, -1, 1), horizon=4):
    """穷举全部候选；目标仅用预测，不访问未来环境真值。"""
    best = None
    for index, sequence in enumerate(product(actions, repeat=horizon)):
        imagined = rollout(model, previous, current, sequence)
        velocity = (imagined[-1] - imagined[-2]) / DT
        cost = ((imagined[-1] - goal) ** 2 + 0.1 * velocity ** 2
                + 0.01 * sum(a * a for a in sequence))
        key = (round(cost, 12), index)  # 仅稳定近似相等的候选次序。
        if best is None or key < best[0]:
            best = (key, sequence, imagined, cost)
    return {"actions": list(best[1]), "imagined": best[2], "cost": best[3],
            "candidate_count": len(actions) ** horizon}


def control(model, case, actions=(0, -1, 1)):
    x, velocity, goal = case["x"], case["velocity"], case["goal"]
    previous = x - DT * velocity
    trace = []
    for t in range(24):
        plan = choose(model, previous, x, goal, actions)
        action = plan["actions"][0]
        new_x, new_velocity = step(x, velocity, action)
        trace.append({"t": t, "previous": previous, "x": x, "velocity": velocity,
                      "action": action, "next_x": new_x, "next_velocity": new_velocity,
                      "plan": plan})
        previous, x, velocity = x, new_x, new_velocity
    return {"case": case, "trace": trace, "final_error": abs(x - goal),
            "final_speed": abs(velocity),
            "success": abs(x - goal) < 0.15 and abs(velocity) < 0.2}


def image_toy():
    """两个等概率未来；每个未来只有一个 2×2 亮块。"""
    frames = []
    for column in (10, 22):
        frame = [[0.0] * 32 for _ in range(32)]
        for y in (15, 16):
            for x in (column, column + 1):
                frame[y][x] = 1.0
        frames.append(frame)
    mean = [[(frames[0][y][x] + frames[1][y][x]) / 2
             for x in range(32)] for y in range(32)]
    blank = [[0.0] * 32 for _ in range(32)]
    def expected_mse(prediction):
        return sum(sum((prediction[y][x] - actual[y][x]) ** 2
                       for y in range(32) for x in range(32)) / 1024
                   for actual in frames) / 2
    return {"futures": frames, "mean": mean,
            "expected_mse": {"mean": expected_mse(mean),
                             "one_valid_future": expected_mse(frames[0]),
                             "blank": expected_mse(blank)}}


def mixed_mask():
    labels = ["文0", "BOI", "图0", "图1", "EOI",
              "文1", "BOI", "图2", "图3", "EOI"]
    blocks = [None, None, 0, 0, None, None, None, 1, 1, None]
    mask = [[int(j <= i or (blocks[i] is not None and blocks[i] == blocks[j]))
             for j in range(len(labels))] for i in range(len(labels))]
    return {"labels": labels, "allowed": mask}


def shared_mechanism():
    """手工构造的两条冲突预测；不是已训练视频模型的输出。"""
    initial_velocity = 1.0
    actions, predicted_velocities = [1.0, -1.0], [1.2, 0.1]
    separate = [(v - 0.4 * a) / initial_velocity
                for a, v in zip(actions, predicted_velocities)]
    shared = sum(separate) / len(separate)
    error = sum((shared * initial_velocity + 0.4 * a - v) ** 2
                for a, v in zip(actions, predicted_velocities)) / 2
    return {"initial_velocity": initial_velocity, "actions": actions,
            "predictions": predicted_velocities, "separate_retentions": separate,
            "shared_retention": shared, "separate_mse": 0.0, "shared_mse": error}


def small_experiments():
    return {"image": image_toy(), "mask": mixed_mask(), "shared": shared_mechanism()}


def summary(record):
    rows = []
    for regime in ("standard", "strong_commands", "changed_retention"):
        rows.append(f"{regime}：单步 / 第 16 步位置 RMSE（米）")
        for mode, name in zip(MODES, NAMES):
            values = [r["evaluation"][regime][mode] for r in record["runs"]]
            one = sum(v["one_step_rmse"] for v in values) / len(values)
            last = sum(v["rollout_rmse"][-1] for v in values) / len(values)
            rows.append(f"  {name}：{one:.4f} / {last:.4f}")
    rows.append("闭环控制成功次数（3 份拟合 × 18 个固定案例）：")
    for mode in record["control_methods"]:
        cases = [c for run in record["runs"] for c in run["control"][mode]]
        rows.append(f"  {mode}：{sum(c['success'] for c in cases)}/{len(cases)}")
    toy = small_experiments()
    rows.append(f"图像均值的期望 MSE：{toy['image']['expected_mse']['mean']:.6f}")
    rows.append(f"共享惯性系数的拟合 MSE：{toy['shared']['shared_mse']:.4f}")
    return "\n".join(rows)


def main():
    parser = ArgumentParser(description=__doc__)
    parser.add_argument("--small", type=Path)
    args = parser.parse_args()
    if args.small:
        save_new(args.small, small_experiments())
        print(args.small)
        return
    record = json.loads(RECORD.read_text())
    assert record["complete"]
    print(summary(record))


if __name__ == "__main__":
    main()

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

1. 识别检查当前内容；动作条件预测还要处理时间、隐藏状态和指定动作的后果。需要另外评价动作响应和连续预测。
2. 图0、图1 是同一张正在去噪的图片；图2 属于后面另一张图。文0 不能读取未来图0，否则会破坏文字的因果顺序。
3. 八个位置分别预测 0.5，每处期望误差为 0.25，总 MSE 为 $2/1024$；固定一个未来时，一半全对、一半八处错 1，总期望为 $4/1024$。低 MSE 的平均图却有两个半亮块。
4. 速度衰减为 ±0.4 米/秒，位置变化为 ±0.08 米，分别是 2.08 米和 1.92 米。
5. 先用相邻位置差除以 0.2 得当前速度，再代入速度与位置更新。要求 $a_t\in[-1,1]$，否则应使用饱和后的 $u_t$。
6. 不能。开放滚动会把自己的预测送回输入，误差和输入分布都可能改变；单步评价有中间真实观测的帮助。
7. 实际执行量为 1，下一速度为 0.4 米/秒，位置为 0.08 米；线性模型错误地使用 3，预测为 0.24 米。
8. 这一行已经直接使用真实环境规则，失败可能来自有限时域、离散动作、代价权重、固定步数与成功门槛。修改它们时要保留原方案对照，不能只改门槛让分数变好。
9. 共享参数增加了约束，最小误差不可能低于允许每条分支各选参数的结果。差值为零也可能是大家一起拟合得很差，要同时报告误差本身。
10. $5^4=625$ 条，是原来 81 条的约 7.72 倍。没有复用前缀时，每次决策要做 $625\times4=2500$ 次模型转移，原来为 324 次。仅提高成功数还不够，需同时报告预测失配和额外计算。