Skip to content

Ppo

PPO: a Proximal Policy Optimization agent whose whole training run (num_updates iterations of rollout collection + minibatch SGD) is one jax.lax.scan. PPOHparams holds the knobs; ActorCritic (from navix.agents.models) is the network - swap its encoder for partial-observability. Structurally mirrored by navix.agents.pqn.

Bases: PyTreeNode

The encoder carry (navix.agents.models' encoder contract) that went into the network at this step - i.e. the state produced from o_{t-1} and earlier, before o_t was consumed. () for stateless encoders. For TransformerEncoder it's the raw (context, *frame) window ending at o_{t-1}; storing it here and replaying it in ppo_loss (rather than re-deriving the window) is exact, because the window has no dependence on the encoder's parameters (see that class's docstring).

Bases: HParams

Hyperparameters for PPO. Frozen; .replace(...) for a variant. The per-field defaults are tuned for navix's small gridworlds, not copied from a CartPole reference.

Whether to anneal the learning rate linearly to 0 at the end of training.

Number of environment frames to train for.

PPO clip parameter.

Whether to clip the value loss in the PPO loss.

Entropy coefficient in the total loss.

Lambda parameter of the TD(lambda) return.

Starting learning rate.

Maximum gradient norm for clipping.

Whether to normalise the advantages in the PPO loss.

Number of parallel environments to run.

Number of epochs to train for.

Number of minibatches to split the data into for training.

Number of steps to run in each environment per update.

Value function coefficient in the total loss.

Bases: TrainState

The live encoder carry for the num_envs running environments, threaded across collect_experience calls the same way env_state is; also the pre-step carry for the post-rollout bootstrap in update. () for stateless encoders.