#!/usr/bin/env python3 """Reproduce Lab 3's state-only ridge behavior clone using Python's standard library. Usage: python reproduce_bc.py pusht_export.json python reproduce_bc.py pusht_export.json --model-id CHECKPOINT_ID cat pusht_export.json | python reproduce_bc.py - No packages, GPU, account, network, RGB observations, ACT, or diffusion training are required. This reconstructs the exact educational affine ridge model, not a more capable policy. Physical rollout metrics must still be evaluated in the native PhysX lab; low imitation error does not establish task success. """ import argparse import json import math import sys FEATURE_NAMES = ["block_x", "block_z", "block_sin_angle", "block_cos_angle", "block_vx", "block_vz", "block_angular_velocity", "pusher_x", "pusher_z", "pusher_vx", "pusher_vz", "goal_x", "goal_z", "goal_sin_angle", "goal_cos_angle"] BOUND = 1 - 15 / 256 def features(state): block, goal, pusher = state["block"], state["goal"], state["pusher"] velocity = block.get("velocity", [0, 0]) pusher_velocity = state.get("pusherVelocity", [0, 0]) return [*block["position"], math.sin(block["angle"]), math.cos(block["angle"]), *velocity, block.get("angularVelocity", 0), *pusher, *pusher_velocity, *goal["position"], math.sin(goal["angle"]), math.cos(goal["angle"])] def solve(matrix, targets): """Same small partial-pivot Gaussian elimination used by learning.js.""" n, outputs = len(matrix), len(targets[0]) augmented = [row[:] + target[:] for row, target in zip(matrix, targets)] for column in range(n): pivot = max(range(column, n), key=lambda i: abs(augmented[i][column])) augmented[column], augmented[pivot] = augmented[pivot], augmented[column] divisor = augmented[column][column] if abs(divisor) < 1e-14: raise ValueError("Singular fit: add diverse training demonstrations") for j in range(column, n + outputs): augmented[column][j] /= divisor for i in range(n): if i == column: continue factor = augmented[i][column] for j in range(column, n + outputs): augmented[i][j] -= factor * augmented[column][j] return [row[n:] for row in augmented] def predict(model, raw_features): inputs = [1] + [(v - mu) / sd for v, mu, sd in zip(raw_features, model["mean"], model["scale"])] return [max(-BOUND, min(BOUND, sum(v * w[a] for v, w in zip(inputs, model["weights"])))) for a in range(2)] def fit(rows, ridge): xs, ys = [features(r["observation"]) for r in rows], [r["action"] for r in rows] n, d = len(xs), len(FEATURE_NAMES) mean = [sum(x[j] for x in xs) / n for j in range(d)] scale = [math.sqrt(sum((x[j] - mean[j]) ** 2 for x in xs) / n) for j in range(d)] scale = [1 if sd < 1e-7 else sd for sd in scale] normalized = [[1] + [(v - mean[j]) / scale[j] for j, v in enumerate(x)] for x in xs] matrix = [[0.0] * (d + 1) for _ in range(d + 1)] targets = [[0.0, 0.0] for _ in range(d + 1)] for x, action in zip(normalized, ys): for j in range(d + 1): for k in range(d + 1): matrix[j][k] += x[j] * x[k] / n for a in range(2): targets[j][a] += x[j] * action[a] / n for j in range(1, d + 1): matrix[j][j] += ridge model = {"mean": mean, "scale": scale, "weights": solve(matrix, targets)} mse = sum(sum((p - y) ** 2 for p, y in zip(predict(model, x), action)) for x, action in zip(xs, ys)) / (2 * n) return model, mse, xs def reproduce(bundle, model_id=None): if bundle.get("schema") != "artificialgate.foundations.pusht.v1": raise ValueError("Expected an ArtificialGate Lab 3 export bundle") selected_id = model_id or bundle.get("activeModelId") entry = next((item for item in bundle.get("models", []) if item["id"] == selected_id), None) if not entry: raise ValueError("Active checkpoint missing; train a model or pass --model-id") checkpoint = entry["model"] if checkpoint.get("type") != "standardized-ridge-bc" or checkpoint.get("schemaVersion") != 1: raise ValueError("Unsupported checkpoint type/version") if checkpoint.get("actionSpace") != "absolute-pusher-xz-meters" or checkpoint.get("controlDt") != .05: raise ValueError("Expected absolute world X/Z pusher targets in meters at 20 Hz") if checkpoint.get("featureNames") != FEATURE_NAMES or bundle.get("featureNames", FEATURE_NAMES) != FEATURE_NAMES: raise ValueError("Feature schema mismatch") provenance = checkpoint["trainedOn"] if provenance["split"] != "train": raise ValueError("Checkpoint is not train-only") requested = set(provenance["episodeIds"]) selected = [e for e in bundle["episodes"] if e["id"] in requested] if len(selected) != len(requested) or {e["id"] for e in selected} != requested: raise ValueError("Missing or duplicate training episodes in export") if any(e["split"] != "train" for e in selected): raise ValueError("Checkpoint provenance includes a held-out episode") # Preserve dataset ordering: summation order matches the browser fit. rows = [r for e in selected for r in e["transitions"] if r["source"] in ("expert", "human")] if len(rows) < 20 or len(rows) != provenance["samples"]: raise ValueError("Training sample count does not match saved checkpoint") if any(not all(math.isfinite(v) for v in features(r["observation"]) + r["action"]) for r in rows): raise ValueError("Non-finite training data") model, mse, xs = fit(rows, checkpoint["lambda"]) mean_error = max(abs(a - b) for a, b in zip(model["mean"], checkpoint["mean"])) scale_error = max(abs(a - b) for a, b in zip(model["scale"], checkpoint["scale"])) weight_error = max(abs(a - b) for row, saved in zip(model["weights"], checkpoint["weights"]) for a, b in zip(row, saved)) action_error = max(abs(a - b) for x in xs for a, b in zip(predict(model, x), predict(checkpoint, x))) reproduced = mean_error < 1e-6 and scale_error < 1e-6 and weight_error < 1e-6 and action_error < 1e-6 return {"reproduced": reproduced, "model_id": selected_id, "policy": "state-only standardized affine ridge BC", "training_episodes": len(selected), "training_samples": len(rows), "heldout_used_for_fit": 0, "learner_prefix_rows_excluded": sum(r["source"] not in ("expert", "human") for e in selected for r in e["transitions"]), "max_mean_error": mean_error, "max_scale_error": scale_error, "max_weight_error": weight_error, "max_action_error": action_error, "saved_training_mse": checkpoint["trainingMSE"], "reproduced_training_mse": mse, "physical_success_claim": "None: use held-out native PhysX rollouts to measure success."} def main(): parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) parser.add_argument("bundle", help="Exported JSON bundle, or - for stdin") parser.add_argument("--model-id", help="Checkpoint ID; defaults to bundle.activeModelId") args = parser.parse_args() try: if args.bundle == "-": bundle = json.load(sys.stdin) else: with open(args.bundle, encoding="utf-8") as stream: bundle = json.load(stream) report = reproduce(bundle, args.model_id) print(json.dumps(report, indent=2, allow_nan=False)) return 0 if report["reproduced"] else 1 except (ValueError, KeyError, TypeError, OSError) as error: print(f"Reproduction failed: {error}", file=sys.stderr) return 2 if __name__ == "__main__": raise SystemExit(main())