Skills › Scientific Agent Skills › pufferlib › env_template.py
env_template.py
210 lines · skills/pufferlib/scripts/env_template.py at 5f2bc03
1#!/usr/bin/env python3
2"""Dependency-free synthetic Gymnasium-style environment template.
4This module is intentionally local and synthetic. It does not import PufferLib,
5Gymnasium, environment plug-ins, native extensions, or ROMs. Port the contract
6to a separately reviewed Gymnasium or PufferLib environment only after the
7validator passes.
8"""
10from __future__ import annotations
12import argparse
13import math
14import random
15from dataclasses import dataclass
16from typing import Any
18try:
19 from ._common import UserInputError, bounded_int, emit_json
20except ImportError: # Direct script execution.
21 from _common import UserInputError, bounded_int, emit_json
24@dataclass(frozen=True)
25class DiscreteSpace:
26 """Minimal stand-in for a discrete action space."""
28 n: int
30 @property
31 def shape(self) -> tuple[int, ...]:
32 return ()
34 def contains(self, value: Any) -> bool:
35 return isinstance(value, int) and not isinstance(value, bool) and 0 <= value < self.n
37 def sample(self, rng: random.Random) -> int:
38 return rng.randrange(self.n)
41@dataclass(frozen=True)
42class BoxSpace:
43 """Minimal one-dimensional finite box used by the synthetic environment."""
45 low: float
46 high: float
47 shape: tuple[int, ...]
48 dtype: str = "float32"
50 def contains(self, value: Any) -> bool:
51 if len(self.shape) != 1 or not isinstance(value, (list, tuple)):
52 return False
53 if len(value) != self.shape[0]:
54 return False
55 for item in value:
56 if isinstance(item, bool) or not isinstance(item, (int, float)):
57 return False
58 if not math.isfinite(float(item)) or not self.low <= float(item) <= self.high:
59 return False
60 return True
63class SyntheticGymEnv:
64 """Small deterministic environment implementing Gymnasium's five-tuple API."""
66 metadata = {"render_modes": []}
68 def __init__(self, *, max_steps: int = 16) -> None:
69 if not 1 <= max_steps <= 10_000:
70 raise UserInputError("max_steps must be between 1 and 10000")
71 self.max_steps = max_steps
72 self.observation_space = BoxSpace(-1.0, 1.0, (4,))
73 self.action_space = DiscreteSpace(3)
74 self._rng = random.Random()
75 self._initialized = False
76 self._done = False
77 self._position = 0.0
78 self._target = 0.75
79 self._step_count = 0
80 self._last_action = 0.0
82 def _observation(self) -> list[float]:
83 return [
84 float(self._position),
85 float(self._target),
86 float(self._step_count / self.max_steps),
87 float(self._last_action),
88 ]
90 def reset(
91 self,
92 *,
93 seed: int | None = None,
94 options: dict[str, Any] | None = None,
95 ) -> tuple[list[float], dict[str, Any]]:
96 """Reset state and return ``(observation, info)``."""
97 if seed is not None:
98 self._rng.seed(seed)
99 if options is not None and set(options) - {"position", "target"}:
100 raise UserInputError("reset options may contain only position and target")
102 self._position = self._rng.uniform(-0.5, 0.5)
103 self._target = self._rng.choice((-0.75, 0.75))
104 if options:
105 self._position = float(options.get("position", self._position))
106 self._target = float(options.get("target", self._target))
107 if not -1.0 <= self._position <= 1.0 or not -1.0 <= self._target <= 1.0:
108 raise UserInputError("position and target options must be within [-1, 1]")
110 self._step_count = 0
111 self._last_action = 0.0
112 self._initialized = True
113 self._done = False
114 observation = self._observation()
115 return observation, {"seed": seed, "synthetic": True}
117 def step(
118 self, action: int
119 ) -> tuple[list[float], float, bool, bool, dict[str, Any]]:
120 """Advance one step and return the Gymnasium five-tuple."""
121 if not self._initialized:
122 raise RuntimeError("reset() must be called before step()")
123 if self._done:
124 raise RuntimeError("reset() must be called after termination or truncation")
125 if not self.action_space.contains(action):
126 raise ValueError(f"action {action!r} is outside the action space")
128 movement = (-0.125, 0.0, 0.125)[action]
129 self._position = max(-1.0, min(1.0, self._position + movement))
130 self._last_action = movement / 0.125
131 self._step_count += 1
133 distance = abs(self._target - self._position)
134 terminated = distance <= 0.0625
135 truncated = self._step_count >= self.max_steps and not terminated
136 reward = 1.0 if terminated else -distance
137 self._done = terminated or truncated
138 info = {
139 "distance": float(distance),
140 "episode_step": self._step_count,
141 }
142 return self._observation(), float(reward), terminated, truncated, info
144 def close(self) -> None:
145 self._initialized = False
146 self._done = True
149def run_demo(*, seed: int, steps: int, max_steps: int) -> dict[str, Any]:
150 """Run a bounded deterministic rollout for documentation and smoke tests."""
151 env = SyntheticGymEnv(max_steps=max_steps)
152 action_rng = random.Random(seed + 1)
153 observation, _ = env.reset(seed=seed)
154 total_reward = 0.0
155 resets = 0
156 terminated_count = 0
157 truncated_count = 0
159 for index in range(steps):
160 action = env.action_space.sample(action_rng)
161 observation, reward, terminated, truncated, _ = env.step(action)
162 total_reward += reward
163 terminated_count += int(terminated)
164 truncated_count += int(truncated)
165 if terminated or truncated:
166 resets += 1
167 observation, _ = env.reset(seed=seed + resets + index + 1)
169 env.close()
170 return {
171 "environment": "synthetic",
172 "last_observation": observation,
173 "network_used": False,
174 "resets": resets,
175 "seed": seed,
176 "steps": steps,
177 "terminated": terminated_count,
178 "total_reward": total_reward,
179 "truncated": truncated_count,
180 }
183def build_parser() -> argparse.ArgumentParser:
184 parser = argparse.ArgumentParser(
185 description="Run the dependency-free synthetic environment template."
186 )
187 parser.add_argument("--seed", type=int, default=42)
188 parser.add_argument("--steps", type=int, default=16, help="1..10000")
189 parser.add_argument("--max-steps", type=int, default=16, help="1..10000")
190 parser.add_argument("--compact", action="store_true", help="Emit compact JSON")
191 return parser
194def main(argv: list[str] | None = None) -> int:
195 parser = build_parser()
196 args = parser.parse_args(argv)
197 try:
198 steps = bounded_int(args.steps, name="steps", minimum=1, maximum=10_000)
199 max_steps = bounded_int(
200 args.max_steps, name="max_steps", minimum=1, maximum=10_000
201 )
202 result = run_demo(seed=args.seed, steps=steps, max_steps=max_steps)
203 except (UserInputError, ValueError, RuntimeError) as exc:
204 parser.error(str(exc))
205 emit_json(result, pretty=not args.compact)
206 return 0
209if __name__ == "__main__":
210 raise SystemExit(main())