Skills › Scientific Agent Skills › pufferlib › repro_plan.py
repro_plan.py
177 lines · skills/pufferlib/scripts/repro_plan.py at 5f2bc03
1#!/usr/bin/env python3
2"""Generate a bounded reproducibility and held-out evaluation plan."""
4from __future__ import annotations
6import argparse
7from typing import Any
9try:
10 from ._common import (
11 SOURCE_4_COMMIT,
12 STABLE_SDIST_SHA256,
13 UserInputError,
14 bounded_int,
15 emit_json,
16 validate_slug,
17 )
18 from .validate_plan import PROFILES
19except ImportError: # Direct script execution.
20 from _common import (
21 SOURCE_4_COMMIT,
22 STABLE_SDIST_SHA256,
23 UserInputError,
24 bounded_int,
25 emit_json,
26 validate_slug,
27 )
28 from validate_plan import PROFILES
31def generate_plan(
32 *,
33 profile: str,
34 environment: str,
35 base_seed: int,
36 replicates: int,
37 eval_episodes: int,
38 benchmark_repeats: int,
39) -> dict[str, Any]:
40 """Create non-overlapping train/eval seeds and reporting requirements."""
41 train_seeds = [base_seed + index for index in range(replicates)]
42 eval_seed_base = base_seed + 1_000_000
43 eval_seeds = [eval_seed_base + index for index in range(replicates)]
44 if set(train_seeds) & set(eval_seeds):
45 raise UserInputError("training and evaluation seeds overlap")
47 if profile == "pypi-3.0.0":
48 upstream = {
49 "artifact": "pufferlib-3.0.0.tar.gz",
50 "package": "pufferlib==3.0.0",
51 "python": ">=3.9",
52 "sha256": STABLE_SDIST_SHA256,
53 "warning": (
54 "PyPI provides only an sdist. Its build can download and compile native "
55 "dependencies; audit and sandbox the build before installation."
56 ),
57 }
58 else:
59 upstream = {
60 "commit": SOURCE_4_COMMIT,
61 "package": "pufferlib source 4.0",
62 "python": ">=3.10",
63 "torch": ">=2.9",
64 "warning": (
65 "The historical 4.0 snapshot is not current native 5.0 or a PyPI artifact. Pin the "
66 "commit and use an audited CUDA/CPU build environment."
67 ),
68 }
70 return {
71 "benchmarking": {
72 "aggregate": ["median", "p10", "p90"],
73 "exclude_setup": False,
74 "fixed_workload": True,
75 "record": [
76 "CPU model and logical/physical cores",
77 "GPU model, driver, CUDA, and precision when applicable",
78 "OS, Python, PufferLib, NumPy, Gymnasium, and PyTorch versions",
79 "backend, start method, workers, envs, buffers, and batch size",
80 "warmup, repeats, wall time, agent steps, and reset count",
81 ],
82 "repeats": benchmark_repeats,
83 "warning": "Do not compare SPS across changed workloads or hardware.",
84 },
85 "environment": {
86 "name": environment,
87 "record": [
88 "source URL and immutable revision",
89 "license and asset/ROM rights",
90 "environment and wrapper configuration",
91 "observation/action spaces and dtypes",
92 "termination, truncation, autoreset, and frame-skip semantics",
93 ],
94 },
95 "evaluation": {
96 "checkpoint_selected_without_eval_feedback": True,
97 "deterministic_policy_pass": True,
98 "episodes_per_seed": eval_episodes,
99 "learning_disabled": True,
100 "report_per_seed_and_aggregate": True,
101 "seeds": eval_seeds,
102 "separate_environment_instances": True,
103 "stochastic_policy_pass": True,
104 },
105 "network_used": False,
106 "profile": profile,
107 "provenance": {
108 "checkpoint_sha256_required": True,
109 "dependency_lock_required": True,
110 "record_git_diff": True,
111 "record_source_commit": True,
112 "upstream": upstream,
113 },
114 "schema_version": 1,
115 "training": {
116 "determinism_limitations_recorded": True,
117 "seeds": train_seeds,
118 },
119 }
122def build_parser() -> argparse.ArgumentParser:
123 parser = argparse.ArgumentParser(
124 description=(
125 "Emit a local reproducibility/evaluation plan; no packages, checkpoints, "
126 "environments, GPU, or network services are opened."
127 )
128 )
129 parser.add_argument("--profile", choices=PROFILES, default="pypi-3.0.0")
130 parser.add_argument("--environment", default="synthetic")
131 parser.add_argument("--base-seed", type=int, default=42)
132 parser.add_argument("--replicates", type=int, default=3, help="1..32")
133 parser.add_argument("--eval-episodes", type=int, default=100, help="1..10000")
134 parser.add_argument("--benchmark-repeats", type=int, default=5, help="3..20")
135 parser.add_argument("--compact", action="store_true")
136 return parser
139def main(argv: list[str] | None = None) -> int:
140 parser = build_parser()
141 args = parser.parse_args(argv)
142 try:
143 environment = validate_slug(args.environment, name="environment")
144 base_seed = bounded_int(
145 args.base_seed, name="base_seed", minimum=0, maximum=2**32 - 1_000_033
146 )
147 replicates = bounded_int(
148 args.replicates, name="replicates", minimum=1, maximum=32
149 )
150 eval_episodes = bounded_int(
151 args.eval_episodes,
152 name="eval_episodes",
153 minimum=1,
154 maximum=10_000,
155 )
156 repeats = bounded_int(
157 args.benchmark_repeats,
158 name="benchmark_repeats",
159 minimum=3,
160 maximum=20,
161 )
162 plan = generate_plan(
163 profile=args.profile,
164 environment=environment,
165 base_seed=base_seed,
166 replicates=replicates,
167 eval_episodes=eval_episodes,
168 benchmark_repeats=repeats,
169 )
170 except (UserInputError, ValueError) as exc:
171 parser.error(str(exc))
172 emit_json(plan, pretty=not args.compact)
173 return 0
176if __name__ == "__main__":
177 raise SystemExit(main())