Skip to content

Models

Shared neural-network building blocks used across navix's agents.

Three groups live here: PPO's encoder/actor-critic components (MLPEncoder, ConvEncoder, TransformerEncoder, ActorCritic), Dreamer's world-model components (categorical-latent utilities, the symexp-twohot head, and the RSSM's encoder/decoder/prior/posterior networks) - the reusable pieces navix.agents.dreamer.WorldModel wires together into an RSSM - and PQN's normalized Q-network (QNetwork). See navix.agents.dreamer and navix.agents.pqn's module docstrings for the algorithms these latter components implement.

The encoder contract (carry-based)

Every PPO/PQN feature extractor subclasses Encoder and implements the same two-method interface, so an agent's training loop is written once and works for both fully- and partially-observable settings by swapping only the encoder:

  • initial_carry(obs_shape, dtype=float32) -> carry - the encoder's hidden state at an episode boundary. Encoder's default is stateless ((), ignoring both args); a stateful encoder overrides it. dtype is the observation's own dtype, so a raw-frame carry isn't silently upcast (uint8 pixels -> float32).
  • __call__(carry, obs, is_first) -> (carry, features) - consume one observation, emit the next carry and a feature vector. is_first (a bool, broadcast per batch element) marks obs as the first frame of a fresh episode, so a stateful encoder re-initialises its carry there rather than reading history that belongs to the episode that just ended.

The stateless encoders (MLPEncoder, ConvEncoder, and PQN's QMLPEncoder/QConvEncoder) carry () and ignore is_first: threading a carry through an agent that uses them is inert, and their output is identical to the pre-carry versions. TransformerEncoder is the stateful one (issue #169): its carry is a raw window of the last context frames, so a pomdp observation function's single-frame stream becomes a history-conditioned feature without the agent, the environment, or the observation function changing. (Dreamer's RSSM encoder, SymlogEncoder, is a separate thing - different __call__ shape, no carry.)

Bases: Module

PPO's network: two independent Encoder towers (actor, critic) each followed by a linear head - a categorical policy over action_dim and a scalar value. Swap actor_encoder / critic_encoder to change what the agent sees (e.g. TransformerEncoder for frame history); the training loop is unchanged.

Attributes:

Name Type Description
action_dim int

number of discrete actions (len(env.action_set)).
actor_encoder Encoder

Encoder for the policy tower.
critic_encoder Encoder

Encoder for the value tower. Must produce the same carry shape as actor_encoder (they share one carry).

A single carry, shared by the actor and critic encoders. This assumes the two encoders derive their carry the same way from the same observation stream - true for the encoders here (a stateless (), or TransformerEncoder's raw-frame window, which doesn't depend on the encoder's parameters) - so threading one carry and advancing it once per step is correct and avoids the actor's and critic's windows drifting apart when only one of policy/value runs (as in PPO.collect_experience, which calls policy only). () for the stateless encoders.

Raises ValueError if the actor and critic encoders don't agree on the carry (e.g. a stateless actor with a stateful critic, or two TransformerEncoders with different context) - otherwise the mismatch only surfaces as an opaque shape error deep in a later .apply() trace.

Bases: Encoder

strides=(2, 2) on every layer, not flax's nn.Conv default of 1 (full-resolution, no downsampling) - an RL rollout batch is huge (num_steps * num_envs samples backprop'd through at once), and without downsampling each layer's activations stay at the input's full spatial resolution while channels grow 16->32->64, which blows up to tens of GB for a modest image (e.g. a 56x56 partially- observable render) at a real rollout batch size - confirmed via an actual OOM (single-op 24.5GB allocation) benchmarking navix's own PPO with this encoder on observations.rgb_first_person. Standard strided downsampling (every real CNN vision encoder in RL, e.g. Nature DQN's) keeps each layer's activation footprint bounded instead of constant-times-growing-channels.

Bases: Module

__call__(feat)

Returns a prediction of symlog(obs) directly (not a distribution) - reconstruction loss is plain MSE against symlog(obs), the paper's "symlog MSE" observation head, simpler than the image-specific decoder the official implementation uses since navix's observations here are flat vectors, not pixels.

Bases: Module

Base for the PPO/PQN feature extractors - the carry-based encoder contract (see this module's docstring). Subclasses implement

__call__(carry, obs, is_first) -> (carry, features)

and, if they hold history across steps, override initial_carry. The default here is stateless: an empty carry, so threading it through an agent is inert. Not an nn.Module you instantiate directly.

The agents that store a per-step carry and replay it in their loss (PPO, PQN - "Option 1") rely on the carry being a pure function of the observation stream, independent of the encoder's parameters - as TransformerEncoder's raw-frame window is. An encoder whose carry is a learned state (a GRU/SSM hidden state) breaks that: the stored carry was produced under stale parameters, so replaying it is no longer what the current network would compute. Such an encoder needs the agent to recompute the carry over the sequence in the loss instead.

The carry at an episode boundary. Stateless by default ((), ignoring both args); TransformerEncoder overrides this with a zeroed frame window of shape (context, *obs_shape) and dtype dtype - callers pass the observation's own dtype so the window isn't silently upcast (e.g. uint8 pixels -> float32).

Bases: Encoder

Two tanh Dense layers - the default ActorCritic encoder for a flat (fully-observable / pre-flattened) observation. Stateless (see Encoder); its output is hidden_size-wide.

Bases: Module

Hidden layer size of the posterior network's MLP.

Number of classes per categorical latent variable ("classes" in the official implementation).

Number of independent categorical latent variables ("stoch" in the official implementation).

Bases: Module

Hidden layer size of the prior network's MLP.

Number of classes per categorical latent variable ("classes" in the official implementation).

Number of independent categorical latent variables ("stoch" in the official implementation).

__call__(h)

Returns raw logits, (..., num_latents, num_classes) - unimix is applied by the caller (unimix_categorical), not baked in here, so every caller treats prior and posterior identically.

Bases: Encoder

QNetwork's POMDP (partially-observable pixel) feature extractor: same strided-downsampling conv stack as ConvEncoder (see its docstring for why the stride matters), projected through a Dense/ LayerNorm/ReLU head to match QMLPEncoder's regularization - PQN's LayerNorm-for-stability argument applies to the features QNetwork regresses Q-values from regardless of whether they came from a Dense or Conv stack, so this keeps it rather than dropping it for pixels. Stateless (Encoder's () carry).

Bases: Encoder

QNetwork's default (MDP, fully-observable/flattened) feature extractor: Dense/LayerNorm/ReLU stacked twice. LayerNorm after every hidden layer is not incidental here the way it might be in MLPEncoder above: it's the specific regularizer the PQN paper shows keeps online Q-learning convergent with no replay buffer and no target network (see navix.agents.pqn's module docstring) - so unlike ActorCritic's encoders, QNetwork's own encoders (this and QConvEncoder) always keep it, rather than leaving it out like the shared MLPEncoder/ ConvEncoder do. Stateless (Encoder's () carry).

Bases: Module

The Q-network PQN regresses towards Q(lambda) targets - a pluggable Encoder feature extractor (QMLPEncoder by default for MDP/flattened observations, swappable for QConvEncoder for POMDP/ pixel observations, or TransformerEncoder for frame history - same encoder-swap pattern ActorCritic uses) followed by a linear head over action_dim raw Q-values (no output activation). Threads the encoder carry through like ActorCritic.

The encoder's carry (() for the stateless Q-encoders).

Bases: Module

Size of the GRU's deterministic hidden state ("deter" in the official implementation).

Bases: Module

Dreamer's RSSM observation encoder: a symlog-input MLP mapping a raw observation to a embed_size embedding for the posterior. Named for its distinctive rlax.signed_logp1 (symlog) input transform - not part of the PPO/PQN Encoder carry-contract family above (it has a different __call__ shape and no carry).

Bases: Module

One pre-LN transformer encoder block (self-attention + MLP, each with a residual connection and LayerNorm applied before the sub-layer, not after) - the standard modern choice (GPT-2 onwards) over the original "Attention Is All You Need" post-LN block, which needs a learning-rate warmup schedule to train stably; none of navix's other components add one, so pre-LN avoids relying on it.

Bases: Encoder

Issue #169: a pomdp-mode observation function (rgb_first_person/ categorical_first_person/symbolic_first_person) returns a single current-step frame - no history. A lone frame doesn't disambiguate states that look identical from the agent's current viewpoint but differ in what led there (e.g. which direction something moved before this frame reached it), so a policy conditioned on it only sees a Markovian approximation of the true partially-observable state. This encoder instead attends over a short window of the last context frames, so the feature it hands to ActorCritic is conditioned on a real (if bounded) piece of history.

History lives in the encoder's carry, not in the environment, the observation function, or the agent's training state. The carry is the literal (context, *frame_shape) window of the last context raw observations, oldest first; __call__ rolls the new obs in (or, on is_first, refills the whole window with obs so it never reaches back into the episode that just ended) and returns the updated window unchanged as the next carry. Keeping the carry as raw frames - not per-frame embeddings - is deliberate: the window then has no dependence on the encoder's parameters, so an agent that stores it at collection time and reuses it in its loss (rather than re-deriving it) gets the exact same gradient, not an approximation.

frame_encoder is applied to each of the context frames with shared weights (same submodule instance, called once per frame - context is a static field, so this unrolls at trace time and repeated calls reuse its parameters rather than creating context copies). Its output must already be hidden_size-wide, the same implicit dimension contract ActorCritic places on its actor_encoder/critic_encoder.

A learned positional embedding is added per frame position before attention (plain per-position table, not sinusoidal - the context length is fixed and small). The output is the last token's post-attention feature (the current frame, now contextualised by the ones before it), not a pool over all positions, which would blur the current frame's signal into older, less relevant ones.

The frame window at an episode boundary: context zero frames, (context, *obs_shape). obs_shape is a single observation's shape (no batch axis) - the caller vmaps __call__ over the batch, so it vmaps an initial_carry-shaped leaf the same way. dtype is the observation's own dtype: the window stores raw frames, so keeping it (e.g. uint8 for navix's first-person pixel/symbolic observations) rather than upcasting to float32 keeps Buffer.carry from being context-frames-wide and 4x per element.

Bases: Module

A scalar-valued prediction head (reward, value) implemented as classification over bins evenly-spaced bins in symlog space, with a twohot cross-entropy loss - the "symexp twohot" head from the paper. Far more robust to reward/value outliers than a Gaussian/MSE head, since an extreme target only ever saturates the loss for its two nearest bins, not the whole (unbounded) squared-error term.

Sum of KL(post_i || prior_i) over the num_latents independent categoricals (the second-to-last axis) - each categorical's own KL, summed, matching how DreamerV3 treats the full stochastic latent (num_latents categoricals of num_classes each, "stoch"/"classes" in the official implementation) as one joint distribution for the purposes of the free-nats floor.

Samples a one-hot vector from dist, with a straight-through gradient: the forward value is a genuine (discrete) sample, but the backward gradient flows as if the output were dist.probs directly (sg(onehot - probs) + probs has forward value onehot, since onehot - probs is stop-gradiented, but its Jacobian w.r.t. upstream parameters is probs's). This is what makes the RSSM's own latent trainable end-to-end despite being discrete - distinct from Dreamer's actor action sampling, which deliberately does NOT use this (see navix.agents.dreamer's module docstring).

Builds a distrax.Categorical from logits after mixing in a unimix fraction of uniform probability mass across the last axis - the "unimix" trick (1% by default in the paper): guarantees every class keeps at least unimix / num_classes probability, so neither the KL term nor the entropy can collapse to exactly zero, which keeps the prior from ever fully committing and losing gradient signal.