From f445ca7e6ac297733365567595b7c392f4fb905a Mon Sep 17 00:00:00 2001 From: tactino <18781106300@163.com> Date: Tue, 27 Oct 2026 16:29:23 -0400 Subject: [PATCH] stepledger: a step ledger on each side of the environment-trainer boundary LedgerEnv records what each environment took and gave; LedgerVectorEnv (Gymnasium, NextStep and SameStep autoreset) and LedgerVecEnv (Stable-Baselines3) record what the trainer sent and received; compare / `stepledger compare` finds the first difference per env and names the episode, step and field. Values are cast to the declared spaces and hashed; rewards are kept as exact bits. Tests cover clean sync and async runs and five injected faults (float16 observation, stale reward, time limit sent as termination, changed action, skipped step); CI runs them with SB3 and both examples. --- .github/workflows/ci.yml | 35 +++++ .gitignore | 9 ++ LICENSE | 202 +++++++++++++++++++++++++++ README.md | 150 ++++++++++++++++++++- benchmarks/overhead.py | 98 ++++++++++++++ examples/gymnasium_async.py | 61 +++++++++ examples/sb3_pendulum.py | 44 ++++++ pyproject.toml | 47 +++++++ src/stepledger/__init__.py | 35 +++++ src/stepledger/__main__.py | 5 + src/stepledger/_canon.py | 136 +++++++++++++++++++ src/stepledger/_record.py | 254 ++++++++++++++++++++++++++++++++++ src/stepledger/_version.py | 1 + src/stepledger/cli.py | 38 ++++++ src/stepledger/compare.py | 262 ++++++++++++++++++++++++++++++++++++ src/stepledger/env.py | 64 +++++++++ src/stepledger/sb3.py | 87 ++++++++++++ src/stepledger/vector.py | 108 +++++++++++++++ tests/conftest.py | 13 ++ tests/ledger_testkit.py | 156 +++++++++++++++++++++ tests/test_canon.py | 100 ++++++++++++++ tests/test_clean_runs.py | 48 +++++++ tests/test_cli.py | 40 ++++++ tests/test_compare.py | 139 +++++++++++++++++++ tests/test_faults.py | 88 ++++++++++++ tests/test_sb3.py | 69 ++++++++++ 26 files changed, 2287 insertions(+), 2 deletions(-) create mode 100644 .github/workflows/ci.yml create mode 100644 .gitignore create mode 100644 LICENSE create mode 100644 benchmarks/overhead.py create mode 100644 examples/gymnasium_async.py create mode 100644 examples/sb3_pendulum.py create mode 100644 pyproject.toml create mode 100644 src/stepledger/__init__.py create mode 100644 src/stepledger/__main__.py create mode 100644 src/stepledger/_canon.py create mode 100644 src/stepledger/_record.py create mode 100644 src/stepledger/_version.py create mode 100644 src/stepledger/cli.py create mode 100644 src/stepledger/compare.py create mode 100644 src/stepledger/env.py create mode 100644 src/stepledger/sb3.py create mode 100644 src/stepledger/vector.py create mode 100644 tests/conftest.py create mode 100644 tests/ledger_testkit.py create mode 100644 tests/test_canon.py create mode 100644 tests/test_clean_runs.py create mode 100644 tests/test_cli.py create mode 100644 tests/test_compare.py create mode 100644 tests/test_faults.py create mode 100644 tests/test_sb3.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..62acf4a --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,35 @@ +name: ci + +on: + push: + pull_request: + +jobs: + test: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.11", "3.12"] + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + - name: Install CPU torch + run: python -m pip install torch --index-url https://download.pytorch.org/whl/cpu + - name: Install stepledger with SB3 and test extras + run: python -m pip install -e ".[sb3,test]" + - name: Tests (SB3 tests included) + run: python -m pytest -v + - name: Examples + working-directory: ${{ runner.temp }} + run: | + python "$GITHUB_WORKSPACE/examples/gymnasium_async.py" + set +e + python "$GITHUB_WORKSPACE/examples/gymnasium_async.py" --fault --dir faulty + code=$? + set -e + test "$code" -eq 1 # the injected fault must be reported + python "$GITHUB_WORKSPACE/examples/sb3_pendulum.py" --dir sb3-ledger + stepledger compare sb3-ledger sb3-ledger/trainer.npz diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..54eb297 --- /dev/null +++ b/.gitignore @@ -0,0 +1,9 @@ +__pycache__/ +*.py[cod] +*.egg-info/ +.venv/ +.pytest_cache/ +build/ +dist/ +ledger/ +*.npz diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..d645695 --- /dev/null +++ b/LICENSE @@ -0,0 +1,202 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/README.md b/README.md index 28b8dcb..1c68186 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,148 @@ -# step-ledger -Find silent faults at the environment-trainer boundary in RL: a step ledger on each side, compared step by step. Gymnasium and Stable-Baselines3. +# stepledger + +Find silent faults at the environment-trainer boundary in reinforcement learning. + +When an environment runs outside the trainer (in another process or on another machine, behind a vector env or a network bridge), a value can change on its way across: an observation rounded to float16, a stale reward, a time limit sent as a termination, a skipped step whose result is filled in. Nothing raises an error, and the policy usually still learns, so training curves do not show it. stepledger keeps a record on each side of the boundary, compares the two step by step, and tells you the env, episode, step and field of the first difference, without needing a fault-free reference run. + +``` +env 1, episode 2, step 12: reward differs (env gave -11.843963, trainer received -14.121007) +env 3, episode 3, step 7: obs differs (env gave [0.7979722, -0.6026943, -2.675495], trainer received [0.79785156, -0.60253906, -2.6757812]) +env 4, episode 0, step 25: terminated, truncated differ (terminated: env gave False, trainer received True; truncated: env gave True, trainer received False) +3 of 8 envs differ +``` + +## Install + +From a checkout of this repository: + +``` +pip install . # numpy and gymnasium only +pip install ".[sb3]" # adds stable-baselines3 +``` + +Python 3.10 or newer, Gymnasium 1.0 or newer. The package imports as `stepledger`; its distribution name is `plugrl-stepledger`. It is not on PyPI yet, and the PyPI project called `stepledger` is an unrelated package. + +## Quickstart: Stable-Baselines3 + +```python +import functools +import gymnasium as gym +from stable_baselines3 import PPO +from stable_baselines3.common.vec_env import SubprocVecEnv +from stepledger import LedgerEnv +from stepledger.sb3 import LedgerVecEnv + +def make_env(i): + return LedgerEnv(gym.make("Pendulum-v1"), "ledger", index=i) # env side: outermost, in the worker + +if __name__ == "__main__": + venv = SubprocVecEnv([functools.partial(make_env, i) for i in range(4)]) + venv = LedgerVecEnv(venv, "ledger/trainer.npz") # trainer side: innermost + PPO("MlpPolicy", venv).learn(10_000) + venv.close() # both sides write their ledgers +``` + +Then: + +``` +stepledger compare ledger ledger/trainer.npz +``` + +## Quickstart: Gymnasium + +```python +import functools +import gymnasium as gym +from gymnasium.vector import AsyncVectorEnv +from stepledger import LedgerEnv, LedgerVectorEnv, compare + +def make_env(i): + return LedgerEnv(gym.make("CartPole-v1"), "ledger", index=i) # env side: outermost, in the worker + +if __name__ == "__main__": + envs = AsyncVectorEnv([functools.partial(make_env, i) for i in range(4)]) + envs = LedgerVectorEnv(envs, "ledger/trainer.npz") # trainer side: innermost + envs.reset(seed=0) + for _ in range(1000): + envs.step(envs.action_space.sample()) + envs.close() # both sides write their ledgers + print(compare("ledger", "ledger/trainer.npz")) +``` + +Longer versions are in [`examples/`](examples): `gymnasium_async.py` (run it with `--fault` to see a report) and `sb3_pendulum.py`. + +## Where the wrappers go + +The comparison checks exactly what lies between the two wrappers, so their places decide what is checked. + +- `LedgerEnv(env, directory, index)` goes **outermost in the environment's own process**: after your TimeLimit, Monitor and other env-side wrappers, as the last thing before the env is handed to the vector env, the worker or the bridge server. Whatever is inside it counts as the environment; whatever is outside it counts as the boundary. `index` must be the env's position in the trainer's vector env. Each one writes `env-{index}.npz` into `directory` when the env is closed (also on `flush()`, and at interpreter exit if it was never closed). +- `LedgerVectorEnv(vector_env, path)` (Gymnasium) or `stepledger.sb3.LedgerVecEnv(venv, path)` (SB3) goes **innermost on the trainer side**: directly around `AsyncVectorEnv`, `SyncVectorEnv`, `SubprocVecEnv`, `DummyVecEnv` or your bridge client, and inside `VecNormalize`, `VecFrameStack`, `NormalizeObservation` and the like, because those change values on purpose after the boundary. It writes one file with every env's record when closed. +- Both sides must use the same `reward_dtype` (default `np.float32`). `compare` checks that the two sides recorded with the same reward dtype and the same observation and action spaces, and refuses to compare otherwise. + +With `SyncVectorEnv` or `DummyVecEnv` both sides live in one process; the boundary is then the vector env itself (its buffers, dtype casts and autoreset logic), which is still worth checking. + +## What is recorded and compared + +Each side records, per env, a stream of events: `reset(obs)` and `step(action, obs, reward, terminated, truncated)`. In a step event `obs` is what the env returned at that step; at an episode end that is the final observation, not the one from the following reset. End flags are stored as `terminated` and `truncated and not terminated`. + +The trainer side follows how the vector env resets envs: + +- Gymnasium `NextStep` autoreset (the default): the step after an episode end is a reset; it is recorded as a reset, and the action sent for it, which the env ignores, is not recorded. +- Gymnasium `SameStep` autoreset: on an episode end it records the step with `info["final_obs"][i]`, then the reset. +- Gymnasium `Disabled` autoreset is refused. +- SB3: on `done` it records the step with `infos[i]["terminal_observation"]` and `truncated = infos[i]["TimeLimit.truncated"]`, then the reset. Actions are taken in `step_async`. + +Before comparing, values are cast to their declared space: Box, MultiDiscrete and MultiBinary to the space dtype, Discrete to int64, Dict and Tuple leaf by leaf with dict keys sorted. Observations and actions are stored as 64-bit blake2b digests of the cast value (its shape, dtype and bytes). Rewards are cast to `reward_dtype` and stored as their raw bits, which is exact and lets the report print reward values without keeping anything else. Pass `keep_values=True` on both sides to also store observations and actions, so the report prints both sides' values at the first difference. + +`compare` reports, per env, the first event where the two streams differ and which fields differ there: `kind` (one side has a reset where the other has a step), `action`, `obs`, `reward`, `terminated`, `truncated`. Episodes are counted from 0 at the first reset; step 0 is the reset and step k the k-th step of the episode, counted on the env side. It also reports streams of different length and envs present on one side only. + +`stepledger compare ENV_DIR TRAINER_FILE` (or `python -m stepledger compare ...`) prints the report and exits with 0 when there is no difference, 1 when there is one, and 2 when it cannot compare (missing files, different settings on the two sides). + +## What it does not catch + +- Faults inside the environment, that is, inside `LedgerEnv`. Both sides see the same wrong value. +- Differences below the precision of the declared dtypes. An observation declared float32 but produced as float64 is compared after the cast to float32. Rewards are compared as float32 by default, because SB3 stores rewards as float32; use `reward_dtype=np.float64` on both sides with trainers that keep float64 rewards. +- Anything not in the event stream: `info` dicts, seeds, reset options, render frames. +- Changes made by wrappers outside the trainer-side ledger. Those are on the trainer side of the comparison by construction. +- Two different values with the same 64-bit digest. The chance per compared field is about 1 in 2^64. + +## Cost + +Measured with `benchmarks/overhead.py`: a `SyncVectorEnv` of 16 `Pendulum-v1` envs, random actions, 5000 vector steps, median of 5 interleaved runs (i7-1360P laptop, Windows 11, Python 3.12, NumPy 2.5, Gymnasium 1.4). + +| ledgers | time per vector step | vs. none | +| --- | --- | --- | +| none | 180 µs | 1.00x | +| `LedgerEnv` on each env | 258 µs | 1.43x | +| `LedgerVectorEnv` | 240 µs | 1.33x | +| both | 325 µs | 1.81x | + +That is about 5 µs per env step on the env side and 4 µs on the trainer side. Pendulum takes about 11 µs per env step here, so this is close to the worst case in relative terms; for slower environments the ratio is smaller. Hashing time grows with observation size, at about 0.9 GB/s: about 23 µs per side for an 84x84x3 uint8 frame and 170 µs for 224x224x3. + +Storage per event (one reset or one step of one env): + +- in memory while recording: 27 bytes (1 kind, three 8-byte digests, 2 flags), plus the values with `keep_values=True`. The trainer side holds every env, so 16 envs x 1M steps is about 430 MB. +- on disk (compressed `.npz`, episode and step counters included): about 20 bytes on Pendulum; about 35 bytes with `keep_values=True`. + +## Ledger file format + +Each file is an `.npz` without pickles. It holds a JSON `meta` string (format version, side, reward dtype, space descriptions, package version), the env `index` of each stream and stream `offsets`, and these columns, concatenated over streams: `kind` (uint8, 0 reset, 1 step), `h_obs`, `h_act`, `h_rew` (uint64), `terminated`, `truncated` (bool), `episode`, `step` (int64). With `keep_values=True` it also holds `values_obs_{j}` and `values_act_{j}`, one array per space leaf. `stepledger.read(path)` loads one (`.meta`, and `.streams[index].columns`). + +## Citation + +```bibtex +@inproceedings{gou2026silent, + title={Detecting Silent Faults at the Environment-Trainer Boundary in Reinforcement Learning}, + author={Gou, Zuo and Lu, Chenhao}, + booktitle={CoRL 2026 Workshop on Everything Beneath the Policy}, + year={2026} +} +``` + +## Who + +stepledger is part of [PlugRL](https://github.com/PlugRL). It was built by [Chenhao Lu](https://github.com/CTP314), [Zuo Gou](https://github.com/tactino) and [Zilin Kang](https://github.com/nothingbutbut). + +## License + +Apache-2.0 diff --git a/benchmarks/overhead.py b/benchmarks/overhead.py new file mode 100644 index 0000000..d6be882 --- /dev/null +++ b/benchmarks/overhead.py @@ -0,0 +1,98 @@ +"""Time per vector step with and without ledgers, and ledger size per event. + +SyncVectorEnv of 16 Pendulum-v1 envs (default 200-step time limit), random +actions sampled up front. Each setting runs ``--repeats`` times, interleaved; +the median is reported. + + python benchmarks/overhead.py +""" + +from __future__ import annotations + +import argparse +import functools +import platform +import statistics +import tempfile +import time +from pathlib import Path + +import gymnasium as gym +import numpy as np +from gymnasium.vector import SyncVectorEnv + +from stepledger import LedgerEnv, LedgerVectorEnv +from stepledger._record import read + +SETTINGS = { + "no ledgers": (False, False), + "env side only (LedgerEnv x16)": (True, False), + "trainer side only (LedgerVectorEnv)": (False, True), + "both sides": (True, True), +} + + +def make_env(directory: Path, index: int, env_side: bool, keep_values: bool) -> gym.Env: + env = gym.make("Pendulum-v1") + return LedgerEnv(env, directory, index, keep_values=keep_values) if env_side else env + + +def run(directory: Path, env_side: bool, trainer_side: bool, args: argparse.Namespace, keep_values: bool = False) -> float: + fns = [functools.partial(make_env, directory, i, env_side, keep_values) for i in range(args.num_envs)] + vec = SyncVectorEnv(fns) + if trainer_side: + vec = LedgerVectorEnv(vec, directory / "trainer.npz", keep_values=keep_values) + vec.action_space.seed(0) + actions = [vec.action_space.sample() for _ in range(args.steps)] + vec.reset(seed=0) + for action in actions[:200]: # warm-up + vec.step(action) + start = time.perf_counter() + for action in actions: + vec.step(action) + elapsed = time.perf_counter() - start + vec.close() + return elapsed / args.steps + + +def bytes_per_event(directory: Path) -> tuple[float, float]: + """(env side, trainer side) file bytes per recorded event.""" + env_bytes = env_events = 0 + for path in directory.glob("env-*.npz"): + env_bytes += path.stat().st_size + env_events += sum(len(s) for s in read(path).streams.values()) + trainer = directory / "trainer.npz" + trainer_events = sum(len(s) for s in read(trainer).streams.values()) + return env_bytes / env_events, trainer.stat().st_size / trainer_events + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--num-envs", type=int, default=16) + parser.add_argument("--steps", type=int, default=5000) + parser.add_argument("--repeats", type=int, default=5) + args = parser.parse_args() + + print(f"python {platform.python_version()}, numpy {np.__version__}, gymnasium {gym.__version__}, {platform.processor()}") + times: dict[str, list[float]] = {name: [] for name in SETTINGS} + with tempfile.TemporaryDirectory() as tmp: + for r in range(args.repeats): + for name, (env_side, trainer_side) in SETTINGS.items(): + times[name].append(run(Path(tmp) / f"{r}-{name[:3]}", env_side, trainer_side, args)) + base = statistics.median(times["no ledgers"]) + print(f"\n{args.num_envs} envs, {args.steps} vector steps, median of {args.repeats}") + for name, values in times.items(): + t = statistics.median(values) + print(f" {name:38s} {t * 1e6:8.0f} us/vector step {t / base:5.2f}x (+{(t - base) / args.num_envs * 1e6:.1f} us per env step)") + + both = Path(tmp) / "0-bot" + env_b, trainer_b = bytes_per_event(both) + print(f"\nfile size per event: env side {env_b:.1f} B, trainer side {trainer_b:.1f} B") + kept = Path(tmp) / "kept" + run(kept, True, True, args, keep_values=True) + env_b, trainer_b = bytes_per_event(kept) + print(f"with keep_values=True: env side {env_b:.1f} B, trainer side {trainer_b:.1f} B") + + +if __name__ == "__main__": + main() diff --git a/examples/gymnasium_async.py b/examples/gymnasium_async.py new file mode 100644 index 0000000..4e5e771 --- /dev/null +++ b/examples/gymnasium_async.py @@ -0,0 +1,61 @@ +"""Step ledgers on both sides of a Gymnasium AsyncVectorEnv, then compare. + + python examples/gymnasium_async.py # clean run: exit code 0 + python examples/gymnasium_async.py --fault # env 2 hands over one rounded observation: exit code 1 + stepledger compare ledger ledger/trainer.npz +""" + +import argparse +import functools + +import gymnasium as gym +import numpy as np +from gymnasium.vector import AsyncVectorEnv + +import stepledger + + +class RoundOnce(gym.Wrapper): + """Stands in for a boundary that rounds the observation through float16, once, at step 40.""" + + def __init__(self, env: gym.Env) -> None: + super().__init__(env) + self.steps = 0 + + def step(self, action): + obs, reward, terminated, truncated, info = self.env.step(action) + self.steps += 1 + if self.steps == 40: + obs = obs.astype(np.float16).astype(np.float32) + return obs, reward, terminated, truncated, info + + +def make_env(index: int, directory: str, fault: bool) -> gym.Env: + env = gym.make("Pendulum-v1") + env = stepledger.LedgerEnv(env, directory, index) # outermost wrapper in the worker process + if fault and index == 2: + env = RoundOnce(env) # between the env-side ledger and the trainer + return env + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--fault", action="store_true") + parser.add_argument("--dir", default="ledger") + args = parser.parse_args() + + envs = AsyncVectorEnv([functools.partial(make_env, i, args.dir, args.fault) for i in range(4)]) + envs = stepledger.LedgerVectorEnv(envs, f"{args.dir}/trainer.npz") # innermost on the trainer side + envs.action_space.seed(0) + envs.reset(seed=0) + for _ in range(500): + envs.step(envs.action_space.sample()) + envs.close() # the workers write env-0.npz ... env-3.npz, the wrapper writes trainer.npz + + report = stepledger.compare(args.dir, f"{args.dir}/trainer.npz") + print(report) + raise SystemExit(0 if report.clean else 1) + + +if __name__ == "__main__": + main() diff --git a/examples/sb3_pendulum.py b/examples/sb3_pendulum.py new file mode 100644 index 0000000..0800f90 --- /dev/null +++ b/examples/sb3_pendulum.py @@ -0,0 +1,44 @@ +"""Train PPO on Pendulum with step ledgers on both sides of a SubprocVecEnv, then compare. + + pip install ".[sb3]" # from a checkout of this repository + python examples/sb3_pendulum.py + stepledger compare ledger ledger/trainer.npz +""" + +import argparse +import functools + +import gymnasium as gym +from stable_baselines3 import PPO +from stable_baselines3.common.monitor import Monitor +from stable_baselines3.common.vec_env import SubprocVecEnv, VecNormalize + +import stepledger +from stepledger.sb3 import LedgerVecEnv + + +def make_env(index: int, directory: str) -> gym.Env: + env = Monitor(gym.make("Pendulum-v1")) + return stepledger.LedgerEnv(env, directory, index) # outermost wrapper in the worker process + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--dir", default="ledger") + parser.add_argument("--timesteps", type=int, default=8192) + args = parser.parse_args() + + venv = SubprocVecEnv([functools.partial(make_env, i, args.dir) for i in range(4)]) + venv = LedgerVecEnv(venv, f"{args.dir}/trainer.npz") # innermost on the trainer side + venv = VecNormalize(venv) # trainer-side transforms go outside the ledger + model = PPO("MlpPolicy", venv, n_steps=512, seed=0, verbose=0) + model.learn(total_timesteps=args.timesteps) + venv.close() # the workers write env-0.npz ... env-3.npz, LedgerVecEnv writes trainer.npz + + report = stepledger.compare(args.dir, f"{args.dir}/trainer.npz") + print(report) + raise SystemExit(0 if report.clean else 1) + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..10870f8 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,47 @@ +[build-system] +requires = ["hatchling>=1.27"] +build-backend = "hatchling.build" + +[project] +name = "plugrl-stepledger" +dynamic = ["version"] +description = "Find silent faults at the environment-trainer boundary in RL by comparing step ledgers kept on both sides." +readme = "README.md" +license = "Apache-2.0" +license-files = ["LICENSE"] +requires-python = ">=3.10" +authors = [ + { name = "Chenhao Lu" }, + { name = "Zuo Gou" }, + { name = "Zilin Kang" }, +] +keywords = ["reinforcement learning", "gymnasium", "stable-baselines3", "debugging", "vector environments"] +classifiers = [ + "Programming Language :: Python :: 3", + "Operating System :: OS Independent", + "Intended Audience :: Science/Research", + "Topic :: Scientific/Engineering :: Artificial Intelligence", +] +dependencies = [ + "numpy>=1.21", + "gymnasium>=1.0", +] + +[project.optional-dependencies] +sb3 = ["stable-baselines3>=2.5"] +test = ["pytest>=7"] + +[project.scripts] +stepledger = "stepledger.cli:main" + +[project.urls] +Homepage = "https://github.com/PlugRL" + +[tool.hatch.version] +path = "src/stepledger/_version.py" + +[tool.hatch.build.targets.wheel] +packages = ["src/stepledger"] + +[tool.pytest.ini_options] +testpaths = ["tests"] diff --git a/src/stepledger/__init__.py b/src/stepledger/__init__.py new file mode 100644 index 0000000..4736b8b --- /dev/null +++ b/src/stepledger/__init__.py @@ -0,0 +1,35 @@ +"""Find silent faults at the environment-trainer boundary by keeping a step ledger on each side. + +- :class:`LedgerEnv` (env side, outermost wrapper in the env's process) +- :class:`LedgerVectorEnv` (trainer side, innermost around a Gymnasium vector env) +- ``stepledger.sb3.LedgerVecEnv`` (trainer side, innermost around an SB3 VecEnv) +- :func:`compare` or ``stepledger compare ENV_DIR TRAINER_FILE`` +""" + +from stepledger._version import __version__ +from stepledger._record import LedgerFile, read +from stepledger.compare import Difference, EnvResult, MetadataMismatch, Report, compare +from stepledger.env import LedgerEnv +from stepledger.vector import LedgerVectorEnv + +__all__ = [ + "Difference", + "EnvResult", + "LedgerEnv", + "LedgerFile", + "LedgerVectorEnv", + "MetadataMismatch", + "Report", + "__version__", + "compare", + "read", +] + + +def __getattr__(name: str): + # ``stepledger.sb3`` imports stable-baselines3 (and torch), so load it only when asked for. + if name == "sb3": + import importlib + + return importlib.import_module("stepledger.sb3") + raise AttributeError(f"module 'stepledger' has no attribute {name!r}") diff --git a/src/stepledger/__main__.py b/src/stepledger/__main__.py new file mode 100644 index 0000000..9b48e9c --- /dev/null +++ b/src/stepledger/__main__.py @@ -0,0 +1,5 @@ +import sys + +from stepledger.cli import main + +sys.exit(main()) diff --git a/src/stepledger/_canon.py b/src/stepledger/_canon.py new file mode 100644 index 0000000..0813215 --- /dev/null +++ b/src/stepledger/_canon.py @@ -0,0 +1,136 @@ +"""Canonical form and 64-bit digests of observations, actions and rewards. + +A value is canonicalized by casting it to its declared space: every leaf of the +space (a Box, Discrete, MultiDiscrete or MultiBinary inside any nesting of Dict +and Tuple, dict keys sorted) becomes one little-endian numpy array. Both sides +of the boundary canonicalize with the same space, so equal values give equal +digests, and a cast to the declared dtype is never reported as a difference. +""" + +from __future__ import annotations + +import hashlib +import struct +from typing import Any, Sequence + +import numpy as np +from gymnasium import spaces + +KeyPath = tuple # keys from the top-level value down to one leaf + + +def _leaf_dtype(space: spaces.Space) -> np.dtype: + if isinstance(space, spaces.Discrete): + dtype = np.dtype(np.int64) + elif isinstance(space, (spaces.Box, spaces.MultiDiscrete, spaces.MultiBinary)): + dtype = np.dtype(space.dtype) + else: + raise TypeError( + f"stepledger does not support {type(space).__name__} spaces; supported are " + "Box, Discrete, MultiDiscrete, MultiBinary, and Dict/Tuple of these" + ) + return dtype.newbyteorder("<") + + +def _collect(space: spaces.Space, path: KeyPath, out: list[tuple[KeyPath, spaces.Space]]) -> None: + if isinstance(space, spaces.Dict): + for key in sorted(space.spaces, key=str): + _collect(space.spaces[key], path + (key,), out) + elif isinstance(space, spaces.Tuple): + for i, sub in enumerate(space.spaces): + _collect(sub, path + (i,), out) + else: + _leaf_dtype(space) # raises for unsupported spaces + out.append((path, space)) + + +def describe(space: spaces.Space) -> str: + """Describe what the canonical form of ``space`` depends on: types, shapes, dtypes, keys.""" + if isinstance(space, spaces.Dict): + inner = ", ".join(f"{key!r}: {describe(space.spaces[key])}" for key in sorted(space.spaces, key=str)) + return f"Dict({inner})" + if isinstance(space, spaces.Tuple): + return "Tuple(" + ", ".join(describe(sub) for sub in space.spaces) + ")" + return f"{type(space).__name__}(shape={space.shape}, dtype={_leaf_dtype(space).name})" + + +def path_name(path: KeyPath) -> str: + """``('pos', 0)`` -> ``"['pos'][0]"``; the empty path is ``""``.""" + return "".join(f"[{key!r}]" for key in path) + + +class Canon: + """Canonicalizes and digests values of one space.""" + + def __init__(self, space: spaces.Space) -> None: + leaves: list[tuple[KeyPath, spaces.Space]] = [] + _collect(space, (), leaves) + self.space = space + self.paths: list[KeyPath] = [path for path, _ in leaves] + self.dtypes: list[np.dtype] = [_leaf_dtype(sub) for _, sub in leaves] + self.names: list[str] = [path_name(path) for path in self.paths] + self.description = describe(space) + self._single = self.paths == [()] + self._headers: list[dict[tuple, bytes]] = [{} for _ in self.paths] # per leaf: shape -> header + + def leaves(self, value: Any) -> list[np.ndarray]: + """Canonical leaves of one value. + + Also works on a batched value (gymnasium and SB3 batch Dict/Tuple values + leaf by leaf); then row ``i`` of each leaf belongs to env ``i``. + """ + if self._single: + return [np.asarray(value, dtype=self.dtypes[0])] + out = [] + for path, dtype in zip(self.paths, self.dtypes): + item = value + for key in path: + item = item[key] + out.append(np.asarray(item, dtype=dtype)) + return out + + def digest(self, leaves: Sequence[np.ndarray]) -> int: + """64-bit blake2b digest of the canonical leaves. + + For each leaf, in order, the hash is fed its path name (utf-8), its + ndim (uint8) and shape (int64 each, little-endian), its dtype string + (e.g. `` np.dtype: + """Validate a reward dtype: a numeric dtype of at most 8 bytes.""" + dtype = np.dtype(reward_dtype) + if dtype.kind not in "fiub" or dtype.itemsize not in _UINT_OF_SIZE: + raise ValueError(f"reward_dtype must be a numeric dtype of 1, 2, 4 or 8 bytes, got {dtype}") + return dtype.newbyteorder("<") + + +def reward_bits(reward: Any, dtype: np.dtype) -> int: + """The reward cast to ``dtype``, as its raw bits (an exact, reversible stand-in for a digest).""" + value = np.array(reward, dtype=dtype).reshape(()) + return int(value.view(_UINT_OF_SIZE[dtype.itemsize])) + + +def reward_bits_batch(rewards: Any, dtype: np.dtype) -> list[int]: + """:func:`reward_bits` for a batch of rewards.""" + values = np.asarray(rewards, dtype=dtype).reshape(-1) + return values.view(_UINT_OF_SIZE[dtype.itemsize]).tolist() + + +def reward_from_bits(bits: int, dtype: np.dtype) -> np.generic: + """Inverse of :func:`reward_bits`.""" + return np.asarray(bits, dtype=_UINT_OF_SIZE[dtype.itemsize]).view(dtype)[()] diff --git a/src/stepledger/_record.py b/src/stepledger/_record.py new file mode 100644 index 0000000..f9158c5 --- /dev/null +++ b/src/stepledger/_record.py @@ -0,0 +1,254 @@ +"""Event streams, the ledger file format, and flushing open ledgers at exit. + +A ledger file is an ``.npz`` (no pickles) holding one or more event streams, +concatenated column by column: + + meta JSON string: format, side, reward_dtype, spaces, package version + index int64 (k,) env index of each stream + offsets int64 (k+1,) stream s is rows offsets[s]:offsets[s+1] + kind uint8 0 = reset, 1 = step + h_obs uint64 digest of the canonical observation + h_act uint64 digest of the canonical action (0 for resets) + h_rew uint64 raw bits of the reward cast to reward_dtype (0 for resets) + terminated bool + truncated bool canonical: truncated and not terminated + episode int64 0-based episode counter (-1 before the first reset) + step int64 0 for a reset, k for the k-th step of the episode + +With ``keep_values=True`` it also holds ``values_obs_{j}`` and +``values_act_{j}``: leaf ``j`` of the canonical value at every row (zeros where +there is no value, e.g. the action of a reset). +""" + +from __future__ import annotations + +import atexit +import json +import os +import sys +import weakref +from array import array +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Sequence + +import numpy as np +from gymnasium import spaces + +from stepledger._canon import Canon, reward_bits, reward_bits_batch, reward_dtype_of +from stepledger._version import __version__ + +FORMAT = 1 +RESET = 0 +STEP = 1 +COLUMNS = ("kind", "h_obs", "h_act", "h_rew", "terminated", "truncated", "episode", "step") + +Entry = tuple[int, "list[np.ndarray] | None"] # (digest, kept leaves or None) + + +class Stream: + """One sub-environment's events, held in compact arrays (27 bytes per event).""" + + def __init__(self, keep_values: bool) -> None: + self.kind = array("B") + self.h_obs = array("Q") + self.h_act = array("Q") + self.h_rew = array("Q") + self.terminated = array("B") + self.truncated = array("B") + self.obs_values: list[list[np.ndarray]] | None = [] if keep_values else None + self.act_values: list[list[np.ndarray] | None] | None = [] if keep_values else None + + def __len__(self) -> int: + return len(self.kind) + + def add(self, kind: int, obs: Entry, act: Entry | None, rew: int, terminated: bool, truncated: bool) -> None: + self.kind.append(kind) + self.h_obs.append(obs[0]) + self.h_act.append(0 if act is None else act[0]) + self.h_rew.append(rew) + self.terminated.append(terminated) + self.truncated.append(truncated) + if self.obs_values is not None: + self.obs_values.append(obs[1]) + self.act_values.append(None if act is None else act[1]) + + def columns(self) -> dict[str, np.ndarray]: + kind = np.array(self.kind, dtype=np.uint8) + is_reset = kind == RESET + rows = np.arange(len(kind), dtype=np.int64) + last_reset = np.maximum.accumulate(np.where(is_reset, rows, -1)) if len(kind) else rows + return { + "kind": kind, + "h_obs": np.array(self.h_obs, dtype=np.uint64), + "h_act": np.array(self.h_act, dtype=np.uint64), + "h_rew": np.array(self.h_rew, dtype=np.uint64), + "terminated": np.array(self.terminated, dtype=np.uint8).astype(bool), + "truncated": np.array(self.truncated, dtype=np.uint8).astype(bool), + "episode": np.cumsum(is_reset, dtype=np.int64) - 1, + "step": rows - last_reset, + } + + +class Recorder: + """Canonicalizes values and appends events to one stream per sub-environment.""" + + def __init__( + self, + observation_space: spaces.Space, + action_space: spaces.Space, + reward_dtype: Any, + keep_values: bool, + num_streams: int, + ) -> None: + self.obs = Canon(observation_space) + self.act = Canon(action_space) + self.reward_dtype = reward_dtype_of(reward_dtype) + self.keep_values = bool(keep_values) + self.streams = [Stream(self.keep_values) for _ in range(num_streams)] + + def _entry(self, canon: Canon, leaves: Sequence[np.ndarray]) -> Entry: + kept = [np.array(leaf, copy=True) for leaf in leaves] if self.keep_values else None + return canon.digest(leaves), kept + + def action(self, leaves: Sequence[np.ndarray]) -> Entry: + """Digest an action now (before the env can change the array in place).""" + return self._entry(self.act, leaves) + + def reward(self, reward: Any) -> int: + return reward_bits(reward, self.reward_dtype) + + def rewards(self, rewards: Any) -> list[int]: + return reward_bits_batch(rewards, self.reward_dtype) + + def record_reset(self, i: int, obs_leaves: Sequence[np.ndarray]) -> None: + self.streams[i].add(RESET, self._entry(self.obs, obs_leaves), None, 0, False, False) + + def record_step( + self, + i: int, + action: Entry, + obs_leaves: Sequence[np.ndarray], + reward_bits: int, + terminated: Any, + truncated: Any, + ) -> None: + term = bool(terminated) + trunc = bool(truncated) and not term + self.streams[i].add(STEP, self._entry(self.obs, obs_leaves), action, reward_bits, term, trunc) + + def meta(self, **extra: Any) -> dict[str, Any]: + return { + "format": FORMAT, + "stepledger": __version__, + "reward_dtype": self.reward_dtype.name, + "observation_space": self.obs.description, + "action_space": self.act.description, + "keep_values": self.keep_values, + "values": {"obs": self.obs.names, "act": self.act.names} if self.keep_values else None, + **extra, + } + + def write(self, path: Path, indices: Sequence[int], **extra: Any) -> None: + """Write every stream recorded so far to ``path`` (atomically replaced).""" + arrays: dict[str, np.ndarray] = {} + parts = [stream.columns() for stream in self.streams] + lengths = [len(stream) for stream in self.streams] + for name in COLUMNS: + arrays[name] = np.concatenate([part[name] for part in parts]) + arrays["index"] = np.asarray(indices, dtype=np.int64) + arrays["offsets"] = np.concatenate([[0], np.cumsum(lengths)]).astype(np.int64) + arrays["meta"] = np.array(json.dumps(self.meta(**extra))) + if self.keep_values: + obs_rows = [row for stream in self.streams for row in stream.obs_values] + act_rows = [row for stream in self.streams for row in stream.act_values] + for j, dtype in enumerate(self.obs.dtypes): + arrays[f"values_obs_{j}"] = _stack([row[j] for row in obs_rows], dtype) + for j, dtype in enumerate(self.act.dtypes): + arrays[f"values_act_{j}"] = _stack([None if row is None else row[j] for row in act_rows], dtype) + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_name(path.name + ".tmp") + with open(tmp, "wb") as f: + np.savez_compressed(f, **arrays) + os.replace(tmp, path) + + +def _stack(values: list[np.ndarray | None], dtype: np.dtype) -> np.ndarray: + """Stack per-event leaves; missing ones become zeros. Falls back to text if shapes vary.""" + present = [v for v in values if v is not None] + shapes = {v.shape for v in present} + if len(shapes) > 1: + return np.array(["" if v is None else np.array2string(v) for v in values]) + shape = shapes.pop() if shapes else () + out = np.zeros((len(values), *shape), dtype=dtype) + for row, value in enumerate(values): + if value is not None: + out[row] = value + return out + + +@dataclass +class StreamData: + """One stream read back from a ledger file.""" + + index: int + columns: dict[str, np.ndarray] + obs_values: list[np.ndarray] | None # one array per leaf, rows aligned with columns + act_values: list[np.ndarray] | None + + def __len__(self) -> int: + return len(self.columns["kind"]) + + +@dataclass +class LedgerFile: + path: Path + meta: dict[str, Any] + streams: dict[int, StreamData] + + +def read(path: str | os.PathLike) -> LedgerFile: + """Read a ledger file written by :class:`Recorder`.""" + path = Path(path) + with np.load(path, allow_pickle=False) as data: + meta = json.loads(str(data["meta"])) + if meta.get("format") != FORMAT: + raise ValueError(f"{path}: ledger format {meta.get('format')!r}, this stepledger reads format {FORMAT}") + columns = {name: data[name] for name in COLUMNS} + indices = data["index"].tolist() + offsets = data["offsets"].tolist() + values = meta.get("values") or None + obs_values = [data[f"values_obs_{j}"] for j in range(len(values["obs"]))] if values else None + act_values = [data[f"values_act_{j}"] for j in range(len(values["act"]))] if values else None + streams = {} + for s, index in enumerate(indices): + rows = slice(offsets[s], offsets[s + 1]) + streams[index] = StreamData( + index=index, + columns={name: col[rows] for name, col in columns.items()}, + obs_values=None if obs_values is None else [v[rows] for v in obs_values], + act_values=None if act_values is None else [v[rows] for v in act_values], + ) + return LedgerFile(path=path, meta=meta, streams=streams) + + +# Ledgers that have not been closed are flushed when the interpreter exits. +_open_ledgers: "weakref.WeakSet[Any]" = weakref.WeakSet() + + +def track(ledger: Any) -> None: + _open_ledgers.add(ledger) + + +def untrack(ledger: Any) -> None: + _open_ledgers.discard(ledger) + + +@atexit.register +def _flush_open_ledgers() -> None: + for ledger in list(_open_ledgers): + try: + ledger.flush() + except Exception as exc: # never fail interpreter shutdown + print(f"stepledger: could not flush {getattr(ledger, 'path', ledger)}: {exc}", file=sys.stderr) diff --git a/src/stepledger/_version.py b/src/stepledger/_version.py new file mode 100644 index 0000000..3dc1f76 --- /dev/null +++ b/src/stepledger/_version.py @@ -0,0 +1 @@ +__version__ = "0.1.0" diff --git a/src/stepledger/cli.py b/src/stepledger/cli.py new file mode 100644 index 0000000..2277433 --- /dev/null +++ b/src/stepledger/cli.py @@ -0,0 +1,38 @@ +"""Command line: ``stepledger compare ENV_DIR TRAINER_FILE``. + +Exit codes: 0 no differences, 1 differences found, 2 could not compare. +""" + +from __future__ import annotations + +import argparse +import sys + +from stepledger._version import __version__ +from stepledger.compare import compare + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(prog="stepledger", description="Compare env-side and trainer-side step ledgers.") + parser.add_argument("--version", action="version", version=f"stepledger {__version__}") + commands = parser.add_subparsers(dest="command", required=True) + cmd = commands.add_parser( + "compare", + help="find the first difference between the two sides, per env", + description="Exit code 0: no differences. 1: differences found. 2: could not compare.", + ) + cmd.add_argument("env_dir", help="directory holding the env-*.npz ledgers written by LedgerEnv") + cmd.add_argument("trainer_file", help="ledger written by LedgerVectorEnv or sb3.LedgerVecEnv") + args = parser.parse_args(argv) + + try: + report = compare(args.env_dir, args.trainer_file) + except (OSError, ValueError, KeyError) as exc: + print(f"stepledger: {exc}", file=sys.stderr) + return 2 + print(report) + return 0 if report.clean else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/stepledger/compare.py b/src/stepledger/compare.py new file mode 100644 index 0000000..c0198aa --- /dev/null +++ b/src/stepledger/compare.py @@ -0,0 +1,262 @@ +"""Compare env-side ledgers with the trainer-side ledger, event by event.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from pathlib import Path + +import numpy as np + +from stepledger import _record +from stepledger._canon import reward_from_bits + +# (name in reports, column), in the order a step happens: what the trainer sent, then what the env gave. +FIELDS = ( + ("kind", "kind"), + ("action", "h_act"), + ("obs", "h_obs"), + ("reward", "h_rew"), + ("terminated", "terminated"), + ("truncated", "truncated"), +) +# Metadata both sides must agree on, or their digests are not comparable. +MUST_AGREE = ("reward_dtype", "observation_space", "action_space") + + +class MetadataMismatch(ValueError): + """The two sides recorded with different settings, so their ledgers cannot be compared.""" + + +@dataclass(frozen=True) +class Difference: + """The first event at which one env's two streams differ. + + ``event`` is the position in both streams; ``episode`` and ``step`` are + counted on the env side (step 0 is the reset, step k the k-th step). + """ + + index: int + event: int + episode: int + step: int + fields: tuple[str, ...] + details: tuple[str, ...] = () + + def __str__(self) -> str: + verb = "differs" if len(self.fields) == 1 else "differ" + text = f"env {self.index}, episode {self.episode}, step {self.step}: {', '.join(self.fields)} {verb}" + if self.details: + text += " (" + "; ".join(self.details) + ")" + return text + + +@dataclass(frozen=True) +class EnvResult: + """Outcome for one env index. A count is None when that side has no stream for the env.""" + + index: int + env_events: int | None + trainer_events: int | None + difference: Difference | None + + @property + def clean(self) -> bool: + return self.difference is None and self.env_events is not None and self.env_events == self.trainer_events + + def lines(self) -> list[str]: + i, e, t = self.index, self.env_events, self.trainer_events + if e is None: + return [f"env {i}: no env ledger (trainer ledger has {_events(t)})"] + if t is None: + return [f"env {i}: not in the trainer ledger (env ledger has {_events(e)})"] + out = [str(self.difference)] if self.difference else [] + if e != t: + text = f"env {i}: env ledger has {_events(e)}, trainer ledger has {t}" + if self.difference is None: + text += f"; the first {min(e, t)} agree" + out.append(text) + return out + + +@dataclass(frozen=True) +class Report: + """Result of :func:`compare`. ``str(report)`` is a short summary, one line per problem.""" + + envs: tuple[EnvResult, ...] + notes: tuple[str, ...] = () + + @property + def clean(self) -> bool: + return all(env.clean for env in self.envs) + + @property + def differences(self) -> list[Difference]: + return [env.difference for env in self.envs if env.difference is not None] + + def __str__(self) -> str: + lines = list(self.notes) + for env in self.envs: + lines += env.lines() + if self.clean: + events = sum(env.env_events or 0 for env in self.envs) + lines.append(f"no differences: {len(self.envs)} envs, {events} events on each side") + else: + bad = sum(not env.clean for env in self.envs) + lines.append(f"{bad} of {len(self.envs)} envs differ") + return "\n".join(lines) + + +def compare(env_dir: str | os.PathLike, trainer_path: str | os.PathLike) -> Report: + """Compare the ``env-*.npz`` ledgers in ``env_dir`` with the trainer ledger at ``trainer_path``. + + Raises :class:`MetadataMismatch` if the two sides recorded with different + reward dtypes or spaces, and ``FileNotFoundError`` if a side is missing. + """ + env_dir = Path(env_dir) + trainer = _record.read(trainer_path) + _expect_side(trainer, "trainer") + if not env_dir.is_dir(): + raise FileNotFoundError(f"{env_dir} is not a directory") + env_files = sorted(env_dir.glob("env-*.npz")) + if not env_files: + raise FileNotFoundError(f"no env-*.npz ledgers in {env_dir}") + + env_streams: dict[int, _record.StreamData] = {} + versions = set() + for path in env_files: + ledger = _record.read(path) + _expect_side(ledger, "env") + for key in MUST_AGREE: + if ledger.meta[key] != trainer.meta[key]: + raise MetadataMismatch( + f"{key} differs: {ledger.meta[key]!r} in {path.name}, {trainer.meta[key]!r} in " + f"{trainer.path.name}; the two sides must record with the same settings" + ) + for index, stream in ledger.streams.items(): + if index in env_streams: + raise ValueError(f"env {index} has more than one env ledger in {env_dir}") + env_streams[index] = stream + versions.add(ledger.meta["stepledger"]) + + notes = [] + if versions != {trainer.meta["stepledger"]}: + notes.append( + f"note: env ledgers written by stepledger {', '.join(sorted(versions))}, " + f"trainer ledger by {trainer.meta['stepledger']}" + ) + context = _Context( + reward_dtype=np.dtype(trainer.meta["reward_dtype"]), + obs_names=_names(trainer.meta, "obs"), + act_names=_names(trainer.meta, "act"), + ) + indices = sorted(set(env_streams) | set(trainer.streams)) + results = tuple(_compare_env(i, env_streams.get(i), trainer.streams.get(i), context) for i in indices) + return Report(envs=results, notes=tuple(notes)) + + +def _expect_side(ledger: _record.LedgerFile, side: str) -> None: + if ledger.meta.get("side") != side: + raise ValueError(f"{ledger.path} is a {ledger.meta.get('side')}-side ledger, expected {side}-side") + + +def _names(meta: dict, field: str) -> list[str] | None: + values = meta.get("values") + return values[field] if values else None + + +@dataclass(frozen=True) +class _Context: + reward_dtype: np.dtype + obs_names: list[str] | None + act_names: list[str] | None + + +def _compare_env( + index: int, env: _record.StreamData | None, trainer: _record.StreamData | None, context: _Context +) -> EnvResult: + if env is None or trainer is None: + return EnvResult(index, None if env is None else len(env), None if trainer is None else len(trainer), None) + n = min(len(env), len(trainer)) + differs = {} + any_differs = np.zeros(n, dtype=bool) + for field, column in FIELDS: + differs[field] = env.columns[column][:n] != trainer.columns[column][:n] + any_differs |= differs[field] + hits = np.flatnonzero(any_differs) + if hits.size == 0: + return EnvResult(index, len(env), len(trainer), None) + j = int(hits[0]) + fields = tuple(field for field, _ in FIELDS if differs[field][j]) + if "kind" in fields: # a reset against a step: the other fields are not comparable + fields = ("kind",) + details = tuple(_detail(field, j, env, trainer, context, len(fields) > 1) for field in fields) + difference = Difference( + index=index, + event=j, + episode=int(env.columns["episode"][j]), + step=int(env.columns["step"][j]), + fields=fields, + details=tuple(d for d in details if d), + ) + return EnvResult(index, len(env), len(trainer), difference) + + +def _detail( + field: str, j: int, env: _record.StreamData, trainer: _record.StreamData, context: _Context, prefix: bool +) -> str: + """Both sides' values of ``field`` at event ``j``, or "" when they were not kept.""" + if field == "kind": + name = {_record.RESET: "reset", _record.STEP: "step"} + return f"env had a {name[int(env.columns['kind'][j])]}, trainer had a {name[int(trainer.columns['kind'][j])]}" + if field == "reward": + a = reward_from_bits(int(env.columns["h_rew"][j]), context.reward_dtype) + b = reward_from_bits(int(trainer.columns["h_rew"][j]), context.reward_dtype) + text = f"env gave {str(a)}, trainer received {str(b)}" # str(): f-strings print numpy floats as float64 + elif field in ("terminated", "truncated"): + text = f"env gave {bool(env.columns[field][j])}, trainer received {bool(trainer.columns[field][j])}" + else: + names = context.obs_names if field == "obs" else context.act_names + env_values = env.obs_values if field == "obs" else env.act_values + trainer_values = trainer.obs_values if field == "obs" else trainer.act_values + if names is None or env_values is None or trainer_values is None: + return "" + if field == "obs": + text = _leaves_text(names, [v[j] for v in env_values], [v[j] for v in trainer_values], "env gave", "trainer received") + else: + text = _leaves_text(names, [v[j] for v in trainer_values], [v[j] for v in env_values], "trainer sent", "env took") + return f"{field}: {text}" if prefix else text + + +def _leaves_text(names: list[str], first: list[np.ndarray], second: list[np.ndarray], said1: str, said2: str) -> str: + for name, a, b in zip(names, first, second): + if a.shape != b.shape or a.tobytes() != b.tobytes(): + text = _pair(np.asarray(a), np.asarray(b), said1, said2) + return f"{name} {text}" if name else text + return "values equal (digests differ)" + + +def _pair(a: np.ndarray, b: np.ndarray, said1: str, said2: str) -> str: + if a.shape != b.shape: + return f"{said1} shape {a.shape}, {said2} shape {b.shape}" + if a.size <= 8: + return f"{said1} {_show(a)}, {said2} {_show(b)}" + unequal = np.flatnonzero(a.reshape(-1) != b.reshape(-1)) + if unequal.size == 0: + return "values equal (digests differ)" + k = int(unequal[0]) + where = tuple(int(w) for w in np.unravel_index(k, a.shape)) + return f"{unequal.size} of {a.size} elements differ, first at {where}: {said1} {_show(a.flat[k])}, {said2} {_show(b.flat[k])}" + + +def _events(n: int) -> str: + return "1 event" if n == 1 else f"{n} events" + + +def _show(a: np.ndarray | np.generic) -> str: + """Shortest text that tells the value apart in its own dtype (str() of numpy scalars).""" + if a.ndim == 0: + return str(a[()]) + if a.ndim == 1: + return "[" + ", ".join(str(x) for x in a) + "]" + return np.array2string(a, separator=", ") diff --git a/src/stepledger/env.py b/src/stepledger/env.py new file mode 100644 index 0000000..7d23274 --- /dev/null +++ b/src/stepledger/env.py @@ -0,0 +1,64 @@ +"""Env side: record what one environment took and gave.""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any, SupportsFloat + +import gymnasium as gym +import numpy as np + +from stepledger import _record + + +class LedgerEnv(gym.Wrapper): + """Records every reset and step of ``env`` and writes ``env-{index}.npz`` into ``directory``. + + Put it OUTERMOST in the environment's own process: the last wrapper before + the vector env, worker or network bridge, so that what it records is + exactly what crosses the boundary. ``index`` must be the env's position in + the trainer's vector env. + + The file is written on :meth:`close`, on :meth:`flush`, and at interpreter + exit for ledgers that were never closed. + """ + + def __init__( + self, + env: gym.Env, + directory: str | os.PathLike, + index: int, + reward_dtype: Any = np.float32, + keep_values: bool = False, + ) -> None: + super().__init__(env) + self.index = int(index) + self.path = Path(directory) / f"env-{self.index}.npz" + self._recorder = _record.Recorder(env.observation_space, env.action_space, reward_dtype, keep_values, 1) + self._closed = False + _record.track(self) + + def reset(self, *, seed: int | None = None, options: dict[str, Any] | None = None) -> tuple[Any, dict[str, Any]]: + obs, info = self.env.reset(seed=seed, options=options) + self._recorder.record_reset(0, self._recorder.obs.leaves(obs)) + return obs, info + + def step(self, action: Any) -> tuple[Any, SupportsFloat, bool, bool, dict[str, Any]]: + rec = self._recorder + act = rec.action(rec.act.leaves(action)) + obs, reward, terminated, truncated, info = self.env.step(action) + rec.record_step(0, act, rec.obs.leaves(obs), rec.reward(reward), terminated, truncated) + return obs, reward, terminated, truncated, info + + def flush(self) -> Path: + """Write everything recorded so far; returns the file path.""" + self._recorder.write(self.path, [self.index], side="env", index=self.index) + return self.path + + def close(self) -> None: + if not self._closed: + self._closed = True + _record.untrack(self) + self.flush() + super().close() diff --git a/src/stepledger/sb3.py b/src/stepledger/sb3.py new file mode 100644 index 0000000..cf0c2a9 --- /dev/null +++ b/src/stepledger/sb3.py @@ -0,0 +1,87 @@ +"""Trainer side for Stable-Baselines3 vector environments (needs ``stable-baselines3``).""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any + +import numpy as np +from stable_baselines3.common.vec_env import VecEnv, VecEnvWrapper + +from stepledger import _record + + +class LedgerVecEnv(VecEnvWrapper): + """Records what the trainer sent to and received from each sub-environment. + + Put it INNERMOST, directly around ``DummyVecEnv`` / ``SubprocVecEnv`` (or + the bridge client) and inside ``VecNormalize``, ``VecMonitor``, + ``VecFrameStack`` and the like, so it sees the values as they arrive. + + SB3 resets an env as soon as its episode ends: on ``done`` the step is + recorded with ``infos[i]["terminal_observation"]``, then the reset with the + returned observation. ``truncated`` is ``infos[i]["TimeLimit.truncated"]``. + """ + + def __init__( + self, + venv: VecEnv, + path: str | os.PathLike, + reward_dtype: Any = np.float32, + keep_values: bool = False, + ) -> None: + super().__init__(venv) + self.path = Path(path) + self._recorder = _record.Recorder( + venv.observation_space, venv.action_space, reward_dtype, keep_values, venv.num_envs + ) + self._sent: list[_record.Entry] | None = None + self._closed = False + _record.track(self) + + def reset(self) -> Any: + obs = self.venv.reset() + rec = self._recorder + leaves = rec.obs.leaves(obs) + for i in range(self.num_envs): + rec.record_reset(i, [leaf[i] for leaf in leaves]) + self._sent = None + return obs + + def step_async(self, actions: np.ndarray) -> None: + rec = self._recorder + leaves = rec.act.leaves(actions) + self._sent = [rec.action([leaf[i] for leaf in leaves]) for i in range(self.num_envs)] + self.venv.step_async(actions) + + def step_wait(self) -> Any: + obs, rewards, dones, infos = self.venv.step_wait() + if self._sent is None: + raise RuntimeError("LedgerVecEnv.step_wait() called without step_async()") + sent, self._sent = self._sent, None + rec = self._recorder + obs_leaves = rec.obs.leaves(obs) + reward_bits = rec.rewards(rewards) + for i in range(self.num_envs): + row = [leaf[i] for leaf in obs_leaves] + if dones[i]: + truncated = bool(infos[i].get("TimeLimit.truncated", False)) + final_obs = rec.obs.leaves(infos[i]["terminal_observation"]) + rec.record_step(i, sent[i], final_obs, reward_bits[i], not truncated, truncated) + rec.record_reset(i, row) + else: + rec.record_step(i, sent[i], row, reward_bits[i], False, False) + return obs, rewards, dones, infos + + def flush(self) -> Path: + """Write everything recorded so far; returns the file path.""" + self._recorder.write(self.path, list(range(self.num_envs)), side="trainer", trainer="sb3") + return self.path + + def close(self) -> None: + if not self._closed: + self._closed = True + _record.untrack(self) + self.flush() + self.venv.close() diff --git a/src/stepledger/vector.py b/src/stepledger/vector.py new file mode 100644 index 0000000..d327def --- /dev/null +++ b/src/stepledger/vector.py @@ -0,0 +1,108 @@ +"""Trainer side for Gymnasium vector environments.""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any + +import gymnasium as gym +import numpy as np + +from stepledger import _record + +NEXT_STEP = "NextStep" +SAME_STEP = "SameStep" +DISABLED = "Disabled" + + +def autoreset_mode(vector_env: gym.vector.VectorEnv) -> str: + """The vector env's autoreset mode as a string; Gymnasium's default is NextStep.""" + mode = vector_env.metadata.get("autoreset_mode", NEXT_STEP) + return str(getattr(mode, "value", mode)) + + +class LedgerVectorEnv(gym.vector.VectorWrapper): + """Records what the trainer sent to and received from each sub-environment. + + Put it INNERMOST on the trainer side, directly around the vector env (or the + bridge client), so it sees the values as they arrive, before any + normalization or other wrapper. Handles the NextStep and SameStep autoreset + modes; Disabled is refused because the wrapper could not tell which resets + belong to which episode end. + """ + + def __init__( + self, + env: gym.vector.VectorEnv, + path: str | os.PathLike, + reward_dtype: Any = np.float32, + keep_values: bool = False, + ) -> None: + super().__init__(env) + self.mode = autoreset_mode(env) + if self.mode == DISABLED: + raise ValueError( + "LedgerVectorEnv does not support autoreset_mode=Disabled; " + "use NextStep or SameStep (gymnasium.vector.AutoresetMode)" + ) + if self.mode not in (NEXT_STEP, SAME_STEP): + raise ValueError(f"unknown autoreset mode {self.mode!r}") + self.path = Path(path) + self._recorder = _record.Recorder( + env.single_observation_space, env.single_action_space, reward_dtype, keep_values, env.num_envs + ) + self._autoreset = np.zeros(env.num_envs, dtype=bool) + self._closed = False + _record.track(self) + + def reset( + self, *, seed: int | list[int] | None = None, options: dict[str, Any] | None = None + ) -> tuple[Any, dict[str, Any]]: + mask = None + if options is not None and "reset_mask" in options: + mask = np.array(options["reset_mask"], dtype=bool) # read before the env pops it + obs, info = self.env.reset(seed=seed, options=options) + envs = range(self.num_envs) if mask is None else np.flatnonzero(mask).tolist() + rec = self._recorder + leaves = rec.obs.leaves(obs) + for i in envs: + rec.record_reset(i, [leaf[i] for leaf in leaves]) + self._autoreset[i] = False + return obs, info + + def step(self, actions: Any) -> tuple[Any, Any, Any, Any, dict[str, Any]]: + rec = self._recorder + act_leaves = rec.act.leaves(actions) + sent = [rec.action([leaf[i] for leaf in act_leaves]) for i in range(self.num_envs)] + obs, rewards, terminations, truncations, infos = self.env.step(actions) + obs_leaves = rec.obs.leaves(obs) + reward_bits = rec.rewards(rewards) + for i in range(self.num_envs): + row = [leaf[i] for leaf in obs_leaves] + if self.mode == NEXT_STEP and self._autoreset[i]: + # The env was reset; the action sent for it was ignored. + rec.record_reset(i, row) + elif self.mode == SAME_STEP and (terminations[i] or truncations[i]): + final_obs = infos["final_obs"][i] + rec.record_step(i, sent[i], rec.obs.leaves(final_obs), reward_bits[i], terminations[i], truncations[i]) + rec.record_reset(i, row) + else: + rec.record_step(i, sent[i], row, reward_bits[i], terminations[i], truncations[i]) + if self.mode == NEXT_STEP: + self._autoreset = np.logical_or(terminations, truncations) + return obs, rewards, terminations, truncations, infos + + def flush(self) -> Path: + """Write everything recorded so far; returns the file path.""" + self._recorder.write( + self.path, list(range(self.num_envs)), side="trainer", trainer="gymnasium", autoreset_mode=self.mode + ) + return self.path + + def close(self, **kwargs: Any) -> Any: + if not self._closed: + self._closed = True + _record.untrack(self) + self.flush() + return self.env.close(**kwargs) diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..06864f2 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,13 @@ +import sys +from pathlib import Path + +import pytest + +# Make ledger_testkit importable by the tests and by vector-env worker +# processes, whatever pytest import mode is in use. +sys.path.insert(0, str(Path(__file__).parent)) + + +@pytest.fixture +def ledger_dir(tmp_path: Path) -> Path: + return tmp_path / "ledger" diff --git a/tests/ledger_testkit.py b/tests/ledger_testkit.py new file mode 100644 index 0000000..bd4095e --- /dev/null +++ b/tests/ledger_testkit.py @@ -0,0 +1,156 @@ +"""Env factories and fault wrappers shared by the tests. + +This is a plain module (not a conftest) so that vector-env worker processes can +import the factories by name. +""" + +from __future__ import annotations + +import functools +from pathlib import Path +from typing import Any + +import gymnasium as gym +import numpy as np +from gymnasium import spaces +from gymnasium.vector import AsyncVectorEnv, SyncVectorEnv + +from stepledger import LedgerEnv, LedgerVectorEnv, Report, compare +from stepledger._record import read + +# Short time limits so that runs see many episode ends: Pendulum only ever +# truncates; CartPole with random actions mostly terminates, sometimes truncates. +EPISODE_STEPS = {"Pendulum-v1": 25, "CartPole-v1": 30} +FAULTS = ("obs_float16", "stale_reward", "timelimit_as_termination", "action", "skip_step") + + +class Fault(gym.Wrapper): + """Changes what crosses the boundary, once. + + Sits between the env-side LedgerEnv and the vector env. It fires at + ``(episode, step)`` counted the way the trainer sees them (episode 0 starts + at the first reset, step k is the k-th step of the episode), except + ``timelimit_as_termination``, which fires at the first truncation. + """ + + def __init__(self, env: gym.Env, kind: str, episode: int, step: int) -> None: + super().__init__(env) + if kind not in FAULTS: + raise ValueError(kind) + self.kind = kind + self.at = (episode, step) + self.episode = -1 + self.t = 0 + self.last: tuple | None = None + self.fired = False + + def reset(self, **kwargs: Any): + self.episode += 1 + self.t = 0 + return self.env.reset(**kwargs) + + def step(self, action: Any): + self.t += 1 + hit = (self.episode, self.t) == self.at + if hit and self.kind == "skip_step": + return self.last # the previous step's result; the env is not stepped + if hit and self.kind == "action": + action = other_action(self.action_space, action) + obs, reward, terminated, truncated, info = self.env.step(action) + if hit and self.kind == "obs_float16": + obs = round_float16(obs) + if hit and self.kind == "stale_reward": + reward = self.last[1] + if self.kind == "timelimit_as_termination" and truncated and not terminated and not self.fired: + self.fired = True + terminated, truncated = True, False + self.last = (obs, reward, terminated, truncated, info) + return self.last + + +def round_float16(obs: Any) -> Any: + if isinstance(obs, dict): + return {key: round_float16(value) for key, value in obs.items()} + obs = np.asarray(obs) + return obs.astype(np.float16).astype(obs.dtype) + + +def other_action(space: spaces.Space, action: Any) -> Any: + if isinstance(space, spaces.Discrete): + return np.int64(space.start + (int(action) - space.start + 1) % space.n) + return (np.asarray(action) + 0.25).astype(space.dtype) + + +class DictObs(gym.ObservationWrapper): + """Pendulum with a Dict observation, to exercise nested spaces.""" + + def __init__(self, env: gym.Env) -> None: + super().__init__(env) + box = env.observation_space + self.observation_space = spaces.Dict( + { + "velocity": spaces.Box(box.low[2:], box.high[2:], dtype=np.float32), + "angle": spaces.Box(box.low[:2], box.high[:2], dtype=np.float32), + } + ) + + def observation(self, obs: np.ndarray) -> dict: + return {"angle": obs[:2], "velocity": obs[2:]} + + +def make_env( + env_id: str, + directory: str, + index: int, + fault: str | None = None, + at: tuple[int, int] = (1, 7), + fault_index: int = 1, + keep_values: bool = False, + reward_dtype: Any = np.float32, + dict_obs: bool = False, +) -> gym.Env: + env = gym.make(env_id, max_episode_steps=EPISODE_STEPS[env_id]) + if dict_obs: + env = DictObs(env) + env = LedgerEnv(env, directory, index, reward_dtype=reward_dtype, keep_values=keep_values) # outermost in-process + if fault is not None and index == fault_index: + env = Fault(env, fault, *at) # between the env-side ledger and the vector env + return env + + +def env_fns(env_id: str, directory: str | Path, num_envs: int, **kwargs: Any) -> list: + return [functools.partial(make_env, env_id, str(directory), i, **kwargs) for i in range(num_envs)] + + +def run_gymnasium( + directory: Path, + env_id: str = "Pendulum-v1", + vector: str = "sync", + mode: str = "NextStep", + num_envs: int = 3, + steps: int = 100, + seed: int = 0, + trainer_kwargs: dict[str, Any] | None = None, + **env_kwargs: Any, +) -> Report: + """Run random actions through a ledgered Gymnasium vector env and compare the two sides.""" + fns = env_fns(env_id, directory, num_envs, **env_kwargs) + cls = SyncVectorEnv if vector == "sync" else AsyncVectorEnv + vec = LedgerVectorEnv(cls(fns, autoreset_mode=mode), directory / "trainer.npz", **(trainer_kwargs or {})) + vec.action_space.seed(seed) + vec.reset(seed=seed) + for _ in range(steps): + vec.step(vec.action_space.sample()) + vec.close() + return compare(directory, directory / "trainer.npz") + + +def episode_ends(directory: Path) -> tuple[int, int, int]: + """(resets, terminations, truncations) over all env-side ledgers in ``directory``.""" + resets = terminated = truncated = 0 + for path in directory.glob("env-*.npz"): + for stream in read(path).streams.values(): + resets += int(np.sum(stream.columns["kind"] == 0)) + terminated += int(np.sum(stream.columns["terminated"])) + truncated += int(np.sum(stream.columns["truncated"])) + return resets, terminated, truncated diff --git a/tests/test_canon.py b/tests/test_canon.py new file mode 100644 index 0000000..cfa1f19 --- /dev/null +++ b/tests/test_canon.py @@ -0,0 +1,100 @@ +"""Canonical values: the same value digests the same on both sides, a changed one does not.""" + +import hashlib +import struct + +import numpy as np +import pytest +from gymnasium import spaces + +from stepledger._canon import Canon, describe, reward_bits, reward_dtype_of, reward_from_bits + + +def digest(space, value): + canon = Canon(space) + return canon.digest(canon.leaves(value)) + + +def test_box_values_are_cast_to_the_space_dtype(): + box = spaces.Box(-1, 1, (3,), np.float32) + value = np.array([0.1, -0.2, 0.3]) + assert digest(box, value) == digest(box, value.astype(np.float32)) == digest(box, value.tolist()) + assert digest(box, value) != digest(box, value.astype(np.float16)) + assert digest(box, value) != digest(box, value[:2]) + + +def test_shape_is_part_of_the_digest(): + box = spaces.Box(-1, 1, (4,), np.float32) + flat = np.zeros(4, dtype=np.float32) + assert digest(box, flat) != digest(box, flat.reshape(2, 2)) + + +def test_discrete_values_are_int64(): + space = spaces.Discrete(4) + assert digest(space, 2) == digest(space, np.int64(2)) == digest(space, np.int32(2)) == digest(space, np.array(2)) + assert digest(space, 2) != digest(space, 3) + + +def test_multi_discrete_and_multi_binary_use_the_space_dtype(): + md = spaces.MultiDiscrete([3, 4]) + assert digest(md, [1, 2]) == digest(md, np.array([1, 2], dtype=np.int32)) + mb = spaces.MultiBinary(3) + assert Canon(mb).dtypes == [np.dtype(np.int8)] + assert digest(mb, [1, 0, 1]) != digest(mb, [1, 1, 1]) + + +def test_dict_keys_are_sorted_and_tuples_nest(): + space = spaces.Dict( + {"b": spaces.Discrete(2), "a": spaces.Tuple((spaces.Box(0, 1, (2,)), spaces.Discrete(3)))} + ) + canon = Canon(space) + assert canon.names == ["['a'][0]", "['a'][1]", "['b']"] + one = {"b": 1, "a": (np.array([0.5, 0.25]), 2)} + other = {"a": (np.array([0.5, 0.25]), 2), "b": 1} + assert digest(space, one) == digest(space, other) + assert digest(space, one) != digest(space, {"b": 1, "a": (np.array([0.5, 0.25]), 1)}) + assert describe(space) == ( + "Dict('a': Tuple(Box(shape=(2,), dtype=float32), Discrete(shape=(), dtype=int64)), " + "'b': Discrete(shape=(), dtype=int64))" + ) + + +def test_batched_leaves_give_the_same_digest_per_row(): + space = spaces.Dict({"x": spaces.Box(-1, 1, (2,)), "n": spaces.Discrete(5)}) + canon = Canon(space) + batch = {"x": np.array([[0.1, 0.2], [0.3, 0.4]], dtype=np.float32), "n": np.array([3, 4])} + leaves = canon.leaves(batch) + for i in range(2): + single = {"x": batch["x"][i], "n": int(batch["n"][i])} + assert canon.digest([leaf[i] for leaf in leaves]) == digest(space, single) + + +def test_digest_is_the_documented_hash(): + # name + ndim (uint8) + shape (int64 LE) + dtype string + C-order bytes, per leaf + space = spaces.Dict({"b": spaces.Discrete(3), "a": spaces.Box(0, 1, (2, 2), np.float32)}) + a = np.arange(4, dtype=np.float32).reshape(2, 2) + h = hashlib.blake2b(digest_size=8) + h.update(b"['a']" + struct.pack(" 3 * 3 + assert truncated > 0 + if env_id == "CartPole-v1": + assert terminated > 0 + else: + assert terminated == 0 + + +@pytest.mark.parametrize("mode", ["NextStep", "SameStep"]) +def test_clean_run_with_kept_values_and_dict_observations(ledger_dir, mode): + report = run_gymnasium( + ledger_dir, mode=mode, steps=80, keep_values=True, dict_obs=True, trainer_kwargs={"keep_values": True} + ) + assert report.clean, str(report) + + +def test_partial_reset_with_reset_mask(ledger_dir): + vec = LedgerVectorEnv(SyncVectorEnv(env_fns("Pendulum-v1", ledger_dir, 3)), ledger_dir / "trainer.npz") + vec.reset(seed=0) + for t in range(60): + vec.step(vec.action_space.sample()) + if t == 10: + vec.reset(options={"reset_mask": np.array([True, False, True])}) + vec.close() + report = compare(ledger_dir, ledger_dir / "trainer.npz") + assert report.clean, str(report) diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..2278b3d --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,40 @@ +"""The command line: exit 0 when clean, 1 on any difference, 2 when it cannot compare.""" + +import subprocess +import sys + +from ledger_testkit import run_gymnasium + + +def cli(*args): + return subprocess.run([sys.executable, "-m", "stepledger", *map(str, args)], capture_output=True, text=True) + + +def test_exit_code_0_when_clean(ledger_dir): + run_gymnasium(ledger_dir, steps=40) + result = cli("compare", ledger_dir, ledger_dir / "trainer.npz") + assert result.returncode == 0, result.stdout + result.stderr + assert result.stdout.startswith("no differences: 3 envs") + + +def test_exit_code_1_on_a_difference(ledger_dir): + run_gymnasium(ledger_dir, steps=40, fault="stale_reward", at=(1, 7)) + result = cli("compare", ledger_dir, ledger_dir / "trainer.npz") + assert result.returncode == 1, result.stdout + result.stderr + assert result.stdout.startswith("env 1, episode 1, step 7: reward differs (env gave ") + + +def test_exit_code_2_when_it_cannot_compare(tmp_path): + result = cli("compare", tmp_path / "missing", tmp_path / "trainer.npz") + assert result.returncode == 2 + assert result.stderr.startswith("stepledger: ") + + +def test_version(): + result = cli("--version") + assert result.returncode == 0 and result.stdout.startswith("stepledger ") + + +def test_base_package_does_not_import_sb3_or_torch(): + code = "import sys, stepledger; assert 'stable_baselines3' not in sys.modules and 'torch' not in sys.modules" + subprocess.run([sys.executable, "-c", code], check=True) diff --git a/tests/test_compare.py b/tests/test_compare.py new file mode 100644 index 0000000..e917a3b --- /dev/null +++ b/tests/test_compare.py @@ -0,0 +1,139 @@ +"""compare(): metadata checks, envs present on one side only, differing stream lengths, flushing.""" + +import subprocess +import sys +import textwrap + +import gymnasium as gym +import numpy as np +import pytest +from gymnasium import spaces +from gymnasium.vector import AutoresetMode, SyncVectorEnv + +from ledger_testkit import env_fns, make_env, run_gymnasium +from stepledger import LedgerEnv, LedgerVectorEnv, MetadataMismatch, compare +from stepledger._record import read + + +def test_different_reward_dtypes_are_refused(ledger_dir): + with pytest.raises(MetadataMismatch, match="reward_dtype"): + run_gymnasium(ledger_dir, reward_dtype=np.float64) + + +class DeclareFloat16(gym.ObservationWrapper): + """A boundary that changes the declared observation dtype.""" + + def __init__(self, env): + super().__init__(env) + box = env.observation_space + self.observation_space = spaces.Box(box.low.astype(np.float16), box.high.astype(np.float16), dtype=np.float16) + + def observation(self, obs): + return obs.astype(np.float16) + + +def test_different_spaces_are_refused(ledger_dir): + fns = [lambda i=i: DeclareFloat16(make_env("Pendulum-v1", str(ledger_dir), i)) for i in range(2)] + vec = LedgerVectorEnv(SyncVectorEnv(fns), ledger_dir / "trainer.npz") + vec.reset(seed=0) + vec.step(vec.action_space.sample()) + vec.close() + with pytest.raises(MetadataMismatch, match="observation_space"): + compare(ledger_dir, ledger_dir / "trainer.npz") + + +def test_env_on_one_side_only(ledger_dir): + run_gymnasium(ledger_dir, steps=20) + (ledger_dir / "env-2.npz").unlink() + extra = LedgerEnv(gym.make("Pendulum-v1"), ledger_dir, 7) + extra.reset(seed=0) + extra.close() + + report = compare(ledger_dir, ledger_dir / "trainer.npz") + assert not report.clean + assert report.differences == [] + lines = str(report).splitlines() + assert lines[0] == "env 2: no env ledger (trainer ledger has 21 events)" + assert lines[1] == "env 7: not in the trainer ledger (env ledger has 1 event)" + assert lines[-1] == "2 of 4 envs differ" + + +def test_extra_event_on_one_side(ledger_dir): + vec = LedgerVectorEnv(SyncVectorEnv(env_fns("Pendulum-v1", ledger_dir, 2)), ledger_dir / "trainer.npz") + vec.reset(seed=0) + for _ in range(5): + vec.step(vec.action_space.sample()) + vec.env.envs[0].step(np.zeros(1, dtype=np.float32)) # behind the trainer's back + vec.close() + + report = compare(ledger_dir, ledger_dir / "trainer.npz") + assert not report.clean + assert report.differences == [] + assert str(report).splitlines()[0] == "env 0: env ledger has 7 events, trainer ledger has 6; the first 6 agree" + + +def test_disabled_autoreset_is_refused(ledger_dir): + vec = SyncVectorEnv(env_fns("Pendulum-v1", ledger_dir, 2), autoreset_mode=AutoresetMode.DISABLED) + with pytest.raises(ValueError, match="Disabled"): + LedgerVectorEnv(vec, ledger_dir / "trainer.npz") + vec.close() + + +def test_missing_ledgers_and_swapped_arguments(ledger_dir, tmp_path): + run_gymnasium(ledger_dir, steps=5) + with pytest.raises(FileNotFoundError): + compare(tmp_path / "nowhere", ledger_dir / "trainer.npz") + empty = tmp_path / "empty" + empty.mkdir() + with pytest.raises(FileNotFoundError, match="no env-"): + compare(empty, ledger_dir / "trainer.npz") + with pytest.raises(ValueError, match="expected trainer-side"): + compare(ledger_dir, ledger_dir / "env-0.npz") + + +def test_ledger_file_contents(ledger_dir): + run_gymnasium(ledger_dir, num_envs=2, steps=30, mode="SameStep") + env0 = read(ledger_dir / "env-0.npz") + assert env0.meta["side"] == "env" and env0.meta["index"] == 0 + assert env0.meta["reward_dtype"] == "float32" + assert env0.meta["observation_space"] == "Box(shape=(3,), dtype=float32)" + stream = env0.streams[0] + # reset, 25 steps (truncated at the last), reset, 5 steps + assert stream.columns["kind"].tolist() == [0] + [1] * 25 + [0] + [1] * 5 + assert stream.columns["episode"].tolist() == [0] * 26 + [1] * 6 + assert stream.columns["step"].tolist() == list(range(26)) + list(range(6)) + assert stream.columns["truncated"].nonzero()[0].tolist() == [25] + assert not stream.columns["terminated"].any() + trainer = read(ledger_dir / "trainer.npz") + assert trainer.meta["side"] == "trainer" and trainer.meta["autoreset_mode"] == "SameStep" + assert sorted(trainer.streams) == [0, 1] + + +def test_large_values_are_reported_by_their_first_differing_element(): + from stepledger.compare import _pair + + a = np.zeros((4, 4, 3), dtype=np.uint8) + b = a.copy() + b[1, 2, 0] = 7 + b[3, 3, 2] = 1 + assert _pair(a, b, "env gave", "trainer received") == ( + "2 of 48 elements differ, first at (1, 2, 0): env gave 0, trainer received 7" + ) + assert _pair(a, a[:2], "env gave", "trainer received") == "env gave shape (4, 4, 3), trainer received shape (2, 4, 3)" + + +def test_unclosed_ledger_is_written_at_exit(tmp_path): + script = textwrap.dedent( + f""" + import gymnasium as gym + from stepledger import LedgerEnv + env = LedgerEnv(gym.make("Pendulum-v1"), {str(tmp_path)!r}, 4) + env.reset(seed=0) + for _ in range(3): + env.step(env.action_space.sample()) + # no close() + """ + ) + subprocess.run([sys.executable, "-c", script], check=True) + stream = read(tmp_path / "env-4.npz").streams[4] + assert stream.columns["kind"].tolist() == [0, 1, 1, 1] diff --git a/tests/test_faults.py b/tests/test_faults.py new file mode 100644 index 0000000..4d25458 --- /dev/null +++ b/tests/test_faults.py @@ -0,0 +1,88 @@ +"""Faults injected between the env-side ledger and the vector env are found where they happen.""" + +import re + +import pytest + +from ledger_testkit import EPISODE_STEPS, run_gymnasium + +AT = (1, 7) # episode 1, step 7: after the first episode end + +# fault -> (where the first difference must be, the fields it must name) +EXPECTED = { + "obs_float16": (AT, ("obs",)), + "stale_reward": (AT, ("reward",)), + # Pendulum never terminates, so the first truncation is the end of episode 0. + "timelimit_as_termination": ((0, EPISODE_STEPS["Pendulum-v1"]), ("terminated", "truncated")), + # The env took a different action than was sent; what it gave back crossed unchanged. + "action": (AT, ("action",)), + "skip_step": (AT, None), # whatever differs first +} +SETUPS = [("sync", "NextStep"), ("sync", "SameStep"), ("async", "NextStep"), ("async", "SameStep")] + + +@pytest.mark.parametrize("vector,mode", SETUPS) +@pytest.mark.parametrize("fault", sorted(EXPECTED)) +def test_fault_is_found_at_its_env_step_and_field(ledger_dir, fault, vector, mode): + report = run_gymnasium(ledger_dir, "Pendulum-v1", vector, mode, num_envs=3, steps=100, fault=fault, at=AT) + assert not report.clean + assert [d.index for d in report.differences] == [1], str(report) + assert all(env.clean for env in report.envs if env.index != 1) + + where, fields = EXPECTED[fault] + diff = report.differences[0] + assert (diff.episode, diff.step) == where, str(report) + if fields is not None: + assert diff.fields == fields, str(report) + else: # skip_step: a step whose result was filled in, then one event fewer on the env side + assert diff.fields and "kind" not in diff.fields, str(report) + env = report.envs[1] + assert env.env_events == env.trainer_events - 1, str(report) + + +def test_report_line_shows_reward_values(ledger_dir): + report = run_gymnasium(ledger_dir, fault="stale_reward", at=AT) + line = str(report).splitlines()[0] + assert re.fullmatch( + r"env 1, episode 1, step 7: reward differs \(env gave -?[\d.e-]+, trainer received -?[\d.e-]+\)", line + ), line + assert str(report).splitlines()[-1] == "1 of 3 envs differ" + + +def test_kept_values_are_printed_for_observations(ledger_dir): + report = run_gymnasium( + ledger_dir, fault="obs_float16", at=AT, keep_values=True, trainer_kwargs={"keep_values": True} + ) + line = str(report.differences[0]) + assert line.startswith("env 1, episode 1, step 7: obs differs (env gave [") + assert "trainer received [" in line + + +def test_kept_values_name_the_leaf_of_a_dict_observation(ledger_dir): + report = run_gymnasium( + ledger_dir, + mode="SameStep", + fault="obs_float16", + at=AT, + keep_values=True, + dict_obs=True, + trainer_kwargs={"keep_values": True}, + ) + diff = report.differences[0] + assert (diff.index, diff.episode, diff.step, diff.fields) == (1, 1, 7, ("obs",)) + assert diff.details[0].startswith("['angle'] env gave [") + + +def test_changed_discrete_action_is_shown_from_both_sides(ledger_dir): + report = run_gymnasium( + ledger_dir, + "CartPole-v1", + fault="action", + at=(0, 3), + keep_values=True, + trainer_kwargs={"keep_values": True}, + ) + diff = report.differences[0] + assert (diff.index, diff.episode, diff.step) == (1, 0, 3) + assert diff.fields == ("action",) + assert re.fullmatch(r"trainer sent [01], env took [01]", diff.details[0]), diff.details diff --git a/tests/test_sb3.py b/tests/test_sb3.py new file mode 100644 index 0000000..574802d --- /dev/null +++ b/tests/test_sb3.py @@ -0,0 +1,69 @@ +"""Stable-Baselines3 trainer side. Skipped when stable-baselines3 is not installed.""" + +import numpy as np +import pytest + +pytest.importorskip("stable_baselines3") + +from stable_baselines3.common.vec_env import DummyVecEnv, SubprocVecEnv # noqa: E402 + +from ledger_testkit import env_fns, episode_ends # noqa: E402 +from stepledger import compare # noqa: E402 +from stepledger.sb3 import LedgerVecEnv # noqa: E402 + +VEC_ENVS = {"dummy": DummyVecEnv, "subproc": SubprocVecEnv} + + +def run_sb3(directory, env_id, vec, num_envs=3, steps=150, trainer_kwargs=None, **env_kwargs): + venv = LedgerVecEnv( + VEC_ENVS[vec](env_fns(env_id, directory, num_envs, **env_kwargs)), + directory / "trainer.npz", + **(trainer_kwargs or {}), + ) + venv.seed(0) + venv.action_space.seed(0) + venv.reset() + for _ in range(steps): + venv.step(np.stack([venv.action_space.sample() for _ in range(num_envs)])) + venv.close() + return compare(directory, directory / "trainer.npz") + + +@pytest.mark.parametrize("vec", sorted(VEC_ENVS)) +@pytest.mark.parametrize("env_id", ["Pendulum-v1", "CartPole-v1"]) +def test_clean_sb3_run_reports_no_difference(ledger_dir, env_id, vec): + report = run_sb3(ledger_dir, env_id, vec) + assert report.clean, str(report) + resets, terminated, truncated = episode_ends(ledger_dir) + assert resets > 9 and truncated > 0 + assert (terminated > 0) == (env_id == "CartPole-v1") + + +@pytest.mark.parametrize("vec", sorted(VEC_ENVS)) +def test_sb3_stale_reward_is_found(ledger_dir, vec): + report = run_sb3(ledger_dir, "Pendulum-v1", vec, steps=100, fault="stale_reward", at=(1, 7)) + assert [d.index for d in report.differences] == [1], str(report) + diff = report.differences[0] + assert (diff.episode, diff.step, diff.fields) == (1, 7, ("reward",)), str(report) + + +@pytest.mark.parametrize("vec", sorted(VEC_ENVS)) +def test_sb3_time_limit_sent_as_termination_is_found(ledger_dir, vec): + report = run_sb3(ledger_dir, "Pendulum-v1", vec, steps=60, fault="timelimit_as_termination") + diff = report.differences[0] + assert (diff.index, diff.episode, diff.step) == (1, 0, 25), str(report) + assert diff.fields == ("terminated", "truncated"), str(report) + + +def test_sb3_kept_values(ledger_dir): + report = run_sb3( + ledger_dir, + "Pendulum-v1", + "dummy", + steps=60, + fault="obs_float16", + at=(1, 7), + keep_values=True, + trainer_kwargs={"keep_values": True}, + ) + assert str(report.differences[0]).startswith("env 1, episode 1, step 7: obs differs (env gave [")