Skills › Scientific Agent Skills › pufferlib › train_template.py
train_template.py
296 lines · skills/pufferlib/scripts/train_template.py at 5f2bc03
1#!/usr/bin/env python3
2"""Safe PufferLib training-plan template.
4This script never imports PufferLib, starts training, loads checkpoints, uses a
5GPU, or contacts an external logger. It emits a validated argv preview for a
6human to review in an appropriately sandboxed, pinned environment.
7"""
9from __future__ import annotations
11import argparse
12import copy
13from typing import Any
15try:
16 from ._common import LOGGER_CREDENTIAL_ENV, UserInputError, bounded_int, emit_json, validate_slug
17 from .validate_plan import PROFILES, default_plan, validate_plan
18except ImportError: # Direct script execution.
19 from _common import LOGGER_CREDENTIAL_ENV, UserInputError, bounded_int, emit_json, validate_slug
20 from validate_plan import PROFILES, default_plan, validate_plan
23def _command_preview(plan: dict[str, Any]) -> list[str]:
24 """Build argv without shell interpolation or execution."""
25 environment = plan["environment"]["name"]
26 if environment == "synthetic":
27 return []
29 training = plan["training"]
30 vector = plan["vectorization"]
31 command = [
32 "puffer",
33 "train",
34 environment,
35 "--train.total-timesteps",
36 str(training["total_timesteps"]),
37 "--train.seed",
38 str(training["seed"]),
39 "--train.minibatch-size",
40 str(training["minibatch_size"]),
41 ]
42 if plan["profile"] == "pypi-3.0.0":
43 backend_name = {
44 "serial": "Serial",
45 "multiprocessing": "Multiprocessing",
46 "native": "PufferEnv",
47 }[vector["backend"]]
48 command.extend(
49 [
50 "--train.device",
51 training["device"],
52 "--train.bptt-horizon",
53 str(training["horizon"]),
54 "--vec.seed",
55 str(training["seed"]),
56 "--vec.backend",
57 backend_name,
58 "--vec.num-envs",
59 str(vector["num_envs"]),
60 "--vec.num-workers",
61 str(vector["num_workers"]),
62 "--vec.batch-size",
63 str(vector["batch_size"]),
64 ]
65 )
66 else:
67 command.extend(
68 [
69 "--seed",
70 str(training["seed"]),
71 "--train.horizon",
72 str(training["horizon"]),
73 "--vec.total-agents",
74 str(vector["total_agents"]),
75 "--vec.num-buffers",
76 str(vector["num_buffers"]),
77 "--vec.num-threads",
78 str(vector["num_threads"]),
79 ]
80 )
81 if vector["backend"] == "torch":
82 command.append("--slowly")
84 logger = plan["logging"]["backend"]
85 if logger != "none":
86 command.append(f"--{logger}")
87 return command
90def make_plan(args: argparse.Namespace) -> dict[str, Any]:
91 plan = default_plan(args.profile)
92 environment = validate_slug(args.environment, name="environment")
93 plan["environment"]["name"] = environment
94 plan["environment"]["adapter"] = args.adapter
95 plan["environment"]["provenance_verified"] = (
96 environment == "synthetic" or args.provenance_verified
97 )
99 plan["training"].update(
100 {
101 "device": args.device,
102 "horizon": bounded_int(
103 args.horizon, name="horizon", minimum=1, maximum=65_536
104 ),
105 "minibatch_size": bounded_int(
106 args.minibatch_size,
107 name="minibatch_size",
108 minimum=1,
109 maximum=16_777_216,
110 ),
111 "seed": bounded_int(
112 args.seed, name="seed", minimum=0, maximum=2**32 - 1
113 ),
114 "total_timesteps": bounded_int(
115 args.total_timesteps,
116 name="total_timesteps",
117 minimum=1,
118 maximum=1_000_000_000,
119 ),
120 }
121 )
122 plan["evaluation"].update(
123 {
124 "deterministic": args.deterministic_eval,
125 "episodes": bounded_int(
126 args.eval_episodes,
127 name="eval_episodes",
128 minimum=1,
129 maximum=10_000,
130 ),
131 "seed": bounded_int(
132 args.eval_seed, name="eval_seed", minimum=0, maximum=2**32 - 1
133 ),
134 "separate": True,
135 }
136 )
137 plan["logging"].update(
138 {
139 "backend": args.logger,
140 "disclosure_ack": args.acknowledge_external_disclosure,
141 "external_opt_in": args.enable_external_logging,
142 "upload_checkpoints": args.upload_checkpoints,
143 }
144 )
145 if args.profile == "pypi-3.0.0":
146 backend = args.backend or "serial"
147 plan["vectorization"].update(
148 {
149 "backend": backend,
150 "batch_size": bounded_int(
151 args.batch_size,
152 name="batch_size",
153 minimum=1,
154 maximum=65_536,
155 ),
156 "num_envs": bounded_int(
157 args.num_envs, name="num_envs", minimum=1, maximum=65_536
158 ),
159 "num_workers": bounded_int(
160 args.num_workers,
161 name="num_workers",
162 minimum=1,
163 maximum=256,
164 ),
165 "start_method": args.start_method,
166 "zero_copy": args.zero_copy,
167 }
168 )
169 else:
170 backend = args.backend or "torch"
171 plan["vectorization"].update(
172 {
173 "backend": backend,
174 "num_buffers": bounded_int(
175 args.num_buffers, name="num_buffers", minimum=1, maximum=256
176 ),
177 "num_threads": bounded_int(
178 args.num_threads, name="num_threads", minimum=1, maximum=256
179 ),
180 "start_method": "spawn",
181 "total_agents": bounded_int(
182 args.total_agents,
183 name="total_agents",
184 minimum=1,
185 maximum=65_536,
186 ),
187 }
188 )
189 plan["checkpoint"]["format"] = (
190 "state_dict" if backend == "torch" else "native-bin"
191 )
192 return plan
195def build_parser() -> argparse.ArgumentParser:
196 parser = argparse.ArgumentParser(
197 description=(
198 "Create a bounded, local PufferLib dry-run plan. This command never "
199 "executes the generated argv."
200 )
201 )
202 parser.add_argument("--profile", choices=PROFILES, default="pypi-3.0.0")
203 parser.add_argument("--environment", default="synthetic")
204 parser.add_argument(
205 "--adapter",
206 choices=["synthetic", "gymnasium", "pettingzoo", "native-ocean"],
207 default="synthetic",
208 )
209 parser.add_argument(
210 "--provenance-verified",
211 action="store_true",
212 help="Explicitly attest review for a non-synthetic environment",
213 )
214 parser.add_argument(
215 "--backend",
216 choices=["serial", "multiprocessing", "native", "torch"],
217 default=None,
218 )
219 parser.add_argument("--num-envs", type=int, default=4)
220 parser.add_argument("--num-workers", type=int, default=1)
221 parser.add_argument("--batch-size", type=int, default=4)
222 parser.add_argument("--zero-copy", action="store_true")
223 parser.add_argument(
224 "--start-method", choices=["spawn", "forkserver"], default="spawn"
225 )
226 parser.add_argument("--total-agents", type=int, default=64)
227 parser.add_argument("--num-buffers", type=int, default=2)
228 parser.add_argument("--num-threads", type=int, default=1)
229 parser.add_argument("--device", choices=["cpu", "cuda"], default="cpu")
230 parser.add_argument("--total-timesteps", type=int, default=10_000)
231 parser.add_argument("--horizon", type=int, default=16)
232 parser.add_argument("--minibatch-size", type=int, default=256)
233 parser.add_argument("--seed", type=int, default=42)
234 parser.add_argument("--eval-seed", type=int, default=1_000_042)
235 parser.add_argument("--eval-episodes", type=int, default=10)
236 parser.add_argument(
237 "--stochastic-eval",
238 dest="deterministic_eval",
239 action="store_false",
240 default=True,
241 )
242 parser.add_argument("--logger", choices=["none", "wandb", "neptune"], default="none")
243 parser.add_argument("--enable-external-logging", action="store_true")
244 parser.add_argument("--acknowledge-external-disclosure", action="store_true")
245 parser.add_argument("--upload-checkpoints", action="store_true")
246 parser.add_argument("--compact", action="store_true")
247 return parser
250def main(argv: list[str] | None = None) -> int:
251 parser = build_parser()
252 args = parser.parse_args(argv)
253 try:
254 plan = make_plan(args)
255 errors = validate_plan(plan)
256 logger = plan["logging"]["backend"]
257 report = {
258 "command_preview": _command_preview(plan) if not errors else [],
259 "credential": {
260 "environment_variable": LOGGER_CREDENTIAL_ENV.get(logger),
261 "value_read_or_logged": False,
262 },
263 "dry_run": True,
264 "errors": errors,
265 "external_logging_disclosure": (
266 "External services may receive configuration, metrics, source metadata, "
267 "hardware telemetry, and explicitly enabled artifacts; review vendor "
268 "privacy, retention, access, and pricing before use."
269 if logger != "none"
270 else None
271 ),
272 "network_used": False,
273 "preview_limitations": [
274 "Historical 3.0/4.0 CLI preview only; not a 5.0 command.",
275 "Set environment-specific rollout batch size and verify agent counts.",
276 "Start method and zero-copy are plan requirements, not applied by this preview.",
277 "For 4.0, the compiled backend selects CPU/CUDA; the plan does not rebuild it.",
278 ],
279 "plan": copy.deepcopy(plan),
280 "status": "valid" if not errors else "invalid",
281 }
282 except (UserInputError, KeyError, ValueError) as exc:
283 report = {
284 "command_preview": [],
285 "dry_run": True,
286 "errors": [str(exc)],
287 "network_used": False,
288 "plan": None,
289 "status": "invalid",
290 }
291 emit_json(report, pretty=not args.compact)
292 return 0 if report["status"] == "valid" else 1
295if __name__ == "__main__":
296 raise SystemExit(main())