nnx-ppo

Tutorials

  • Quickstart
  • Custom networks 1: composing with containers
  • Custom networks 2: graph networks
  • Custom networks 3: writing your own StatefulModule

Reference

  • Training parameters
  • Logging
  • Checkpointing
  • The rollout_extras channel
  • Batching contract
  • Randomness
  • Containers
  • Pytree utilities
  • Delay and Normalizer placement
  • PPOAdapter reference
  • Recording activations

API Reference

  • Networks
  • Algorithms
  • Wrappers
  • JAX Dataclass Utility
nnx-ppo
  • Overview: module code

All modules for which code is available

  • collections
  • nnx_ppo.algorithms.callbacks
  • nnx_ppo.algorithms.checkpointing
  • nnx_ppo.algorithms.config
  • nnx_ppo.algorithms.metrics
  • nnx_ppo.algorithms.ppo
  • nnx_ppo.algorithms.rollout
  • nnx_ppo.algorithms.types
  • nnx_ppo.jax_dataclass
  • nnx_ppo.networks.adapter
  • nnx_ppo.networks.containers
  • nnx_ppo.networks.delay
  • nnx_ppo.networks.factories
  • nnx_ppo.networks.feedforward
  • nnx_ppo.networks.graph.connection
  • nnx_ppo.networks.graph.graph
  • nnx_ppo.networks.graph.population
  • nnx_ppo.networks.normalizer
  • nnx_ppo.networks.recording
  • nnx_ppo.networks.recurrent
  • nnx_ppo.networks.sampling_layers
  • nnx_ppo.networks.types
  • nnx_ppo.networks.utils
  • nnx_ppo.networks.variational
  • nnx_ppo.wrappers.episode_wrapper
  • nnx_ppo.wrappers.reward_scaling_wrapper
  • typing

© Copyright 2026, Emil Wärnberg.

Built with Sphinx using a theme provided by Read the Docs.