nnx-ppo 0.2.1__tar.gz

nnx_ppo-0.2.1/CHANGELOG.md

@@ -0,0 +1,55 @@
+# Changelog
+
+All notable changes to `nnx-ppo` are recorded here. The format follows
+[Keep a Changelog](https://keepachangelog.com/en/1.1.0/) and the project
+adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
+
+## [0.2.1] — 2026-06-09
+
+### Added
+- `LoggingLevel.THROUGHPUT` — emits `throughput/train_sps`,
+  `throughput/eval_sps`, and `throughput/video_sps` (env + render),
+  with `jax.block_until_ready` barriers so the numbers reflect
+  device-side wall-clock rather than JAX dispatch latency. Included in
+  `LoggingLevel.ALL`.
+- `losses/clipping_fraction` under `LoggingLevel.ACTOR_EXTRA` — the
+  fraction of samples whose likelihood ratio left the PPO clip range
+  during the gradient phase. Tree-mapped, so it works with multi-actor
+  / multi-agent loglikelihoods.
+
+### Changed
+- `RewardScalingWrapper` no longer depends on `mujoco_playground`; it
+  is now typed against the local `RLEnv` / `EnvState` protocols in
+  `nnx_ppo.algorithms.types`.
+- The `[playground]` extra has been removed. `playground` and
+  `warp-lang` are now part of the `[dev]` and `[examples]` extras.
+- Minimum `flax` is now `0.12.7`.
+- License metadata switched to PEP 639 SPDX form
+  (`license = "BSD-3-Clause"` + `license-files = ["LICENSE"]`);
+  requires `setuptools>=77.0` at build time.
+
+### Fixed
+- `PopulationGraph` no longer exposes its build-time registries as a
+  second set of `nnx.Param`s — newer Flax versions reflected through
+  the underscore-prefixed dicts, which tripped `nnx.jit`'s
+  consistent-aliasing check.
+
+### Removed
+- `correlations/action_ll` (under `ACTOR_EXTRA`) — was only emitted
+  for 1-D action spaces and never fired in multi-actuator setups.
+
+## [0.2.0] — 2026-06-03
+
+Initial PyPI release.
+
+### Added
+- Stateful-network PPO training loop (`nnx_ppo.algorithms.ppo.train_ppo`).
+- Network containers (`Sequential`, `Parallel`, `Concat`, `Splitter`) and the
+  two-port `PPOAdapter`.
+- Built-in layers: `Dense`, `LSTM`, `AR1VariationalBottleneck`, `Normalizer`,
+  `Delay`, sampling layers, and graph-population utilities.
+- Rollout machinery with per-environment state reset and `update_statistics`
+  hook for stats-bearing modules.
+- Orbax-based checkpointing.
+- Distillation utility (`nnx_ppo.algorithms.distillation`).
+- Documentation site at <https://nnx-ppo.readthedocs.io>.

nnx_ppo-0.2.1/LICENSE

@@ -0,0 +1,28 @@
+BSD 3-Clause License
+
+Copyright (c) 2026, Emil Wärnberg
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

nnx_ppo-0.2.1/MANIFEST.in

@@ -0,0 +1,12 @@
+include README.md
+include LICENSE
+include CHANGELOG.md
+include pyproject.toml
+include nnx_ppo/py.typed
+recursive-exclude docs/_build *
+recursive-exclude * __pycache__
+recursive-exclude * *.py[cod]
+prune .pytest_cache
+prune wandb
+prune .venv
+prune nnx_ppo.egg-info

nnx_ppo-0.2.1/PKG-INFO

@@ -0,0 +1,136 @@
+Metadata-Version: 2.4
+Name: nnx-ppo
+Version: 0.2.1
+Summary: PPO for stateful/recurrent networks in JAX and flax.nnx.
+Author-email: Emil Wärnberg <ewarnberg@fas.harvard.edu>
+License-Expression: BSD-3-Clause
+Project-URL: Homepage, https://github.com/emiwar/nnx-ppo
+Project-URL: Documentation, https://nnx-ppo.readthedocs.io
+Project-URL: Repository, https://github.com/emiwar/nnx-ppo
+Project-URL: Issues, https://github.com/emiwar/nnx-ppo/issues
+Keywords: reinforcement-learning,ppo,jax,flax,rl,neuroscience
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: jax
+Requires-Dist: flax>=0.12.7
+Requires-Dist: optax
+Requires-Dist: orbax-checkpoint
+Requires-Dist: jaxtyping
+Requires-Dist: numpy
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pyright; extra == "dev"
+Requires-Dist: beartype; extra == "dev"
+Requires-Dist: absl-py; extra == "dev"
+Requires-Dist: playground>=0.2.0; extra == "dev"
+Requires-Dist: warp-lang<1.13; extra == "dev"
+Provides-Extra: examples
+Requires-Dist: brax; extra == "examples"
+Requires-Dist: wandb; extra == "examples"
+Requires-Dist: playground>=0.2.0; extra == "examples"
+Requires-Dist: warp-lang<1.13; extra == "examples"
+Provides-Extra: docs
+Requires-Dist: sphinx>=7.0; extra == "docs"
+Requires-Dist: sphinx-rtd-theme; extra == "docs"
+Dynamic: license-file
+
+# nnx-ppo
+
+Experimental implementation of [Proximal Policy Optimization](https://en.wikipedia.org/wiki/Proximal_policy_optimization)
+in [JAX](https://github.com/google/jax), with first-class support for
+recurrent/stateful networks. Networks are built with
+[flax.nnx](https://flax.readthedocs.io); environments follow the
+[MuJoCo Playground](https://playground.mujoco.org/) API.
+
+> **Status:** experimental — the API may change without notice.
+
+## Highlights
+
+- **Stateful modules.** Recurrent layers, delayed connections, and noisy /
+  variational populations are all first-class citizens. Carry state is
+  threaded through rollout collection *and* multi-epoch loss replay, and is
+  reset correctly when the environment resets — something other JAX RL
+  libraries (e.g. Brax) do not natively support.
+- **PyTree observations.** Observations can be arbitrary nested dicts, which
+  makes it easy to route different streams (proprioception, vision,
+  imitation targets, …) to different parts of a network.
+- **PyTree actions and rewards.** Actions and rewards are also allowed to be
+  PyTrees, which simplifies multi-actuator and multi-agent setups.
+
+## Installation
+
+```bash
+pip install nnx_ppo
+```
+
+`nnx-ppo` installs the CPU build of JAX by default. For a CUDA 12 GPU build:
+
+```bash
+pip install nnx_ppo "jax[cuda12]"
+```
+
+Optional extras:
+
+- `nnx_ppo[examples]` — `brax`, `wandb`, and
+  [`playground`](https://pypi.org/project/playground/) (import name
+  `mujoco_playground`) for the scripts in [examples/](examples/).
+- `nnx_ppo[dev]` — test-suite dependencies (`pytest`, `pyright`, `beartype`,
+  `absl-py`, plus `playground` for the env-driven tests).
+
+## Quick example
+
+```python
+from flax import nnx
+import mujoco_playground
+
+from nnx_ppo.algorithms import ppo
+from nnx_ppo.algorithms.config import TrainConfig, PPOConfig, EvalConfig
+from nnx_ppo.networks.factories import make_mlp_actor_critic
+from nnx_ppo.wrappers import episode_wrapper
+
+env = mujoco_playground.registry.load("CartpoleBalance")
+train_env = episode_wrapper.EpisodeWrapper(env, 1000)
+
+rngs = nnx.Rngs(0)
+nets = make_mlp_actor_critic(
+    env.observation_size,
+    env.action_size,
+    actor_hidden_sizes=[64] * 4,
+    critic_hidden_sizes=[256] * 2,
+    rngs=rngs,
+    normalize_obs=True,
+)
+
+result = ppo.train_ppo(
+    train_env,
+    nets,
+    TrainConfig(
+        ppo=PPOConfig(n_envs=1024, rollout_length=30, total_steps=10_000_000),
+        eval=EvalConfig(enabled=True, every_steps=500_000, n_envs=64,
+                        max_episode_length=1000),
+    ),
+    eval_env=env,
+)
+print(f"Final eval reward: {result.eval_history[-1]['episode_reward_mean']}")
+```
+
+See [examples/wandb_logging.py](examples/wandb_logging.py) for a complete
+training script with W&B logging and video recording.
+
+## Documentation
+
+Full documentation, tutorials, and API reference are at
+<https://nnx-ppo.readthedocs.io>.
+
+## License
+
+BSD 3-Clause — see [LICENSE](LICENSE).

nnx_ppo-0.2.1/README.md

@@ -0,0 +1,91 @@
+# nnx-ppo
+
+Experimental implementation of [Proximal Policy Optimization](https://en.wikipedia.org/wiki/Proximal_policy_optimization)
+in [JAX](https://github.com/google/jax), with first-class support for
+recurrent/stateful networks. Networks are built with
+[flax.nnx](https://flax.readthedocs.io); environments follow the
+[MuJoCo Playground](https://playground.mujoco.org/) API.
+
+> **Status:** experimental — the API may change without notice.
+
+## Highlights
+
+- **Stateful modules.** Recurrent layers, delayed connections, and noisy /
+  variational populations are all first-class citizens. Carry state is
+  threaded through rollout collection *and* multi-epoch loss replay, and is
+  reset correctly when the environment resets — something other JAX RL
+  libraries (e.g. Brax) do not natively support.
+- **PyTree observations.** Observations can be arbitrary nested dicts, which
+  makes it easy to route different streams (proprioception, vision,
+  imitation targets, …) to different parts of a network.
+- **PyTree actions and rewards.** Actions and rewards are also allowed to be
+  PyTrees, which simplifies multi-actuator and multi-agent setups.
+
+## Installation
+
+```bash
+pip install nnx_ppo
+```
+
+`nnx-ppo` installs the CPU build of JAX by default. For a CUDA 12 GPU build:
+
+```bash
+pip install nnx_ppo "jax[cuda12]"
+```
+
+Optional extras:
+
+- `nnx_ppo[examples]` — `brax`, `wandb`, and
+  [`playground`](https://pypi.org/project/playground/) (import name
+  `mujoco_playground`) for the scripts in [examples/](examples/).
+- `nnx_ppo[dev]` — test-suite dependencies (`pytest`, `pyright`, `beartype`,
+  `absl-py`, plus `playground` for the env-driven tests).
+
+## Quick example
+
+```python
+from flax import nnx
+import mujoco_playground
+
+from nnx_ppo.algorithms import ppo
+from nnx_ppo.algorithms.config import TrainConfig, PPOConfig, EvalConfig
+from nnx_ppo.networks.factories import make_mlp_actor_critic
+from nnx_ppo.wrappers import episode_wrapper
+
+env = mujoco_playground.registry.load("CartpoleBalance")
+train_env = episode_wrapper.EpisodeWrapper(env, 1000)
+
+rngs = nnx.Rngs(0)
+nets = make_mlp_actor_critic(
+    env.observation_size,
+    env.action_size,
+    actor_hidden_sizes=[64] * 4,
+    critic_hidden_sizes=[256] * 2,
+    rngs=rngs,
+    normalize_obs=True,
+)
+
+result = ppo.train_ppo(
+    train_env,
+    nets,
+    TrainConfig(
+        ppo=PPOConfig(n_envs=1024, rollout_length=30, total_steps=10_000_000),
+        eval=EvalConfig(enabled=True, every_steps=500_000, n_envs=64,
+                        max_episode_length=1000),
+    ),
+    eval_env=env,
+)
+print(f"Final eval reward: {result.eval_history[-1]['episode_reward_mean']}")
+```
+
+See [examples/wandb_logging.py](examples/wandb_logging.py) for a complete
+training script with W&B logging and video recording.
+
+## Documentation
+
+Full documentation, tutorials, and API reference are at
+<https://nnx-ppo.readthedocs.io>.
+
+## License
+
+BSD 3-Clause — see [LICENSE](LICENSE).

nnx_ppo-0.2.1/nnx_ppo/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.2.1"

nnx_ppo-0.2.1/nnx_ppo/algorithms/callbacks.py

@@ -0,0 +1,35 @@
+"""Callback helpers for training logging."""
+
+from collections.abc import Callable
+
+from nnx_ppo.algorithms.config import VideoData
+
+
+def wandb_video_fn(
+    key: str = "eval_video", fps: int = 30
+) -> Callable[[VideoData], None]:
+    """Create a video callback that logs to wandb.
+
+    Args:
+        key: The wandb log key for the video.
+        fps: Frames per second for the video.
+
+    Returns:
+        A callback function compatible with train_ppo's video_fn parameter.
+
+    Example:
+        >>> result = train_ppo(
+        ...     env, networks,
+        ...     video_fn=wandb_video_fn(fps=50),
+        ... )
+    """
+    import wandb
+
+    def video_fn(data: VideoData) -> None:
+        # Convert THWC to TCHW for wandb
+        video_array = data.frames.transpose(0, 3, 1, 2)
+        wandb.log(
+            {key: wandb.Video(video_array, fps=fps, format="mp4")}, step=data.step
+        )
+
+    return video_fn

nnx_ppo-0.2.1/nnx_ppo/algorithms/checkpointing.py

@@ -0,0 +1,204 @@
+"""Checkpointing utilities for saving and loading training state."""
+
+import os
+import pickle
+from collections.abc import Callable
+from typing import Any, Optional, Protocol, runtime_checkable
+
+import jax
+from flax import nnx
+
+from nnx_ppo.algorithms.config import TrainConfig
+from nnx_ppo.algorithms.types import TrainingState
+
+
+@runtime_checkable
+class CheckpointCallback(Protocol):
+    """Protocol for checkpoint callbacks with named parameters."""
+
+    def __call__(self, training_state: TrainingState, step: int) -> None: ...
+
+
+def _split_net_state(networks):
+    """Split network state: RngKey → pickle, everything else → orbax.
+
+    orbax cannot handle JAX new-style PRNG key arrays (dtype ``key<fry>``), so
+    we separate nnx.RngKey variables and persist them with pickle instead. All
+    other variable types — including nnx.Param, nnx.RngCount, and custom
+    Variable subclasses such as NormalizerStatistics — are saved via orbax.
+
+    Returns:
+        (non_key_state, rng_key_state, abstract_non_key) — the first two are
+        nnx.State objects and the third is the abstract (ShapeDtypeStruct)
+        target needed for orbax restoration.
+    """
+    _, rng_key_state, non_key_state = nnx.split(networks, nnx.RngKey, ...)
+    abstract_non_key = jax.tree_util.tree_map(
+        lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype), non_key_state
+    )
+    return non_key_state, rng_key_state, abstract_non_key
+
+
+def make_checkpoint_fn(
+    directory: str,
+    config: Optional[TrainConfig] = None,
+) -> CheckpointCallback:
+    """Create a checkpoint callback that saves TrainingState to disk.
+
+    Each checkpoint is written to ``{directory}/step_{step:010d}/``, containing:
+
+    - ``networks/`` — orbax checkpoint with all non-PRNG-key network variables
+      (Param, RngCount, NormalizerStatistics, etc.)
+    - ``optimizer/`` — orbax checkpoint with all optimizer state arrays
+    - ``metadata.pkl`` — pickle file with network RngKey variables,
+      all remaining TrainingState fields (``network_states``, ``env_states``,
+      ``rng_key``, ``steps_taken``), the step count, and the optional
+      TrainConfig.
+
+    To resume training from a checkpoint, use :func:`load_checkpoint`.
+
+    Args:
+        directory: Base directory under which checkpoint subdirectories are
+            created.
+        config: Optional TrainConfig to store alongside each checkpoint, useful
+            for reproducing training runs.
+
+    Returns:
+        A callback compatible with train_ppo's ``checkpoint_fn`` parameter.
+
+    Example:
+        >>> result = train_ppo(
+        ...     env, networks, config,
+        ...     checkpoint_fn=make_checkpoint_fn("/tmp/my_run", config=config),
+        ... )
+    """
+
+    abs_directory = os.path.abspath(directory)
+
+    def checkpoint_fn(training_state: TrainingState, step: int) -> None:
+        import orbax.checkpoint as ocp
+
+        step_dir = os.path.join(abs_directory, f"step_{step:010d}")
+        os.makedirs(step_dir, exist_ok=True)
+
+        # Split network state: everything except RngKey → orbax, RngKey → pickle.
+        # orbax cannot handle JAX new-style PRNG key arrays.
+        non_key_state, rng_key_state, _ = _split_net_state(training_state.networks)
+
+        # The optimizer only contains float/int arrays; no key arrays.
+        _, opt_state = nnx.split(training_state.optimizer)
+
+        # Save parameter arrays with orbax. A fresh checkpointer is created per
+        # call and immediately closed to ensure all async writes complete.
+        checkpointer = ocp.StandardCheckpointer()
+        try:
+            checkpointer.save(os.path.join(step_dir, "networks"), non_key_state)
+            checkpointer.save(os.path.join(step_dir, "optimizer"), opt_state)
+        finally:
+            checkpointer.close()
+
+        # Save everything else with pickle (JAX arrays including PRNG keys are
+        # pickle-safe).
+        metadata = {
+            "networks_rng_key_state": rng_key_state,
+            "network_states": training_state.network_states,
+            "env_states": training_state.env_states,
+            "rng_key": training_state.rng_key,
+            "steps_taken": training_state.steps_taken,
+            "step": step,
+            "config": config,
+        }
+        with open(os.path.join(step_dir, "metadata.pkl"), "wb") as f:
+            pickle.dump(metadata, f)
+
+    return checkpoint_fn
+
+
+def load_checkpoint(
+    path: str,
+    networks: Any,
+    optimizer: nnx.Optimizer,
+) -> dict[str, Any]:
+    """Load a checkpoint saved by :func:`make_checkpoint_fn`.
+
+    The ``networks`` and ``optimizer`` arguments serve as structural templates:
+    their architecture must match the checkpoint, but their current parameter
+    values are irrelevant and will be overwritten in-place by the checkpoint
+    values.
+
+    Args:
+        path: Path to the step checkpoint directory, e.g.
+            ``/tmp/my_run/step_0000500000``.
+        networks: Network instance with the same architecture as the checkpoint.
+            Weights are updated in-place.
+        optimizer: Optimizer instance with the same structure as the checkpoint.
+            State is updated in-place.
+
+    Returns:
+        A dict with the following keys:
+
+        - ``"training_state"`` — restored :class:`TrainingState`
+        - ``"step"`` — training step at which the checkpoint was saved (int)
+        - ``"config"`` — :class:`TrainConfig` if one was stored, else ``None``
+
+    Example:
+        >>> networks = factories.make_mlp_actor_critic(...)
+        >>> training_state = ppo.new_training_state(env, networks, n_envs, seed)
+        >>> ckpt = load_checkpoint(
+        ...     "/tmp/my_run/step_0000500000",
+        ...     training_state.networks,
+        ...     training_state.optimizer,
+        ... )
+        >>> result = train_ppo(
+        ...     env, networks, ckpt["config"],
+        ...     initial_state=ckpt["training_state"],
+        ... )
+    """
+    import orbax.checkpoint as ocp
+
+    path = os.path.abspath(path)
+
+    # Build abstract targets from the user-provided templates.
+    # Use ... to capture remaining variables (RngKey) that we restore via pickle.
+    _, _, abstract_non_key = nnx.split(networks, nnx.RngKey, ...)
+    abstract_non_key = jax.tree_util.tree_map(
+        lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype), abstract_non_key
+    )
+    _, opt_template = nnx.split(optimizer)
+    opt_abstract = jax.tree_util.tree_map(
+        lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype), opt_template
+    )
+
+    checkpointer = ocp.StandardCheckpointer()
+    try:
+        restored_non_key = checkpointer.restore(
+            os.path.join(path, "networks"), abstract_non_key
+        )
+        restored_opt = checkpointer.restore(
+            os.path.join(path, "optimizer"), opt_abstract
+        )
+    finally:
+        checkpointer.close()
+
+    with open(os.path.join(path, "metadata.pkl"), "rb") as f:
+        metadata = pickle.load(f)
+
+    # Merge orbax-restored non-key state with pickled rng-key state,
+    # then update the provided modules in-place.
+    full_net_state = nnx.merge_state(restored_non_key, metadata["networks_rng_key_state"])
+    nnx.update(networks, full_net_state)
+    nnx.update(optimizer, restored_opt)
+
+    training_state = TrainingState(
+        networks=networks,
+        network_states=metadata["network_states"],
+        env_states=metadata["env_states"],
+        optimizer=optimizer,
+        rng_key=metadata["rng_key"],
+        steps_taken=metadata["steps_taken"],
+    )
+    return {
+        "training_state": training_state,
+        "step": metadata["step"],
+        "config": metadata["config"],
+    }