jaxdem.rl.models.lstm#

Reinforcement learning model based on a single-layer LSTM.

Classes

LSTMActorCritic(*args, **kwargs)

A recurrent actor–critic with an MLP encoder and an LSTM torso.
class jaxdem.rl.models.lstm.LSTMActorCritic(*args: Any, **kwargs: Any)[source]#

Bases: Model

A recurrent actor–critic with an MLP encoder and an LSTM torso.

This model encodes observations with a small feed-forward network and passes the features through a single-layer LSTM. Linear policy and value heads decode the LSTM hidden state.

Calling modes

  • Sequence mode (training): time-major input x with shape (T, B, obs_dim) produces a distribution and value for every step: policy outputs (T, B, action_space_size) and values (T, B, 1). The carry starts from initial_carry when supplied, otherwise zeros. done[t] clears carry after step t. Sequence evaluation does not overwrite the persistent rollout carry.

  • Single-step mode (evaluation/rollout): input x with shape (..., obs_dim) uses and updates a persistent LSTM carry stored on the module (self.h, self.c); outputs have shape (..., action_space_size) and (..., 1). Use reset to clear state between episodes.

Parameters:

  • observation_space_size (int) – Flattened observation size (obs_dim).

  • action_space_size (int) – Number of action dimensions (for continuous) or number of discrete actions.

  • key (nnx.Rngs) – Random number generator(s) for parameter initialization.

  • hidden_features (int) – Width of the encoder output (and LSTM input).

  • lstm_features (int) – LSTM hidden/state size. Also the feature size consumed by the policy and value heads.

  • activation (Callable) – Activation function applied inside the encoder.

  • action_space (distrax.Bijector | ActionSpace | None) – Bijector to constrain the policy probability distribution (continuous only).

  • cell_type (type[rnn.OptimizedLSTMCell]) – LSTM cell class used for the recurrent layer.

  • remat (bool) – If True, wrap the LSTM scan body with jax.checkpoint to reduce memory in sequence mode.

  • actor_sigma_head (bool) – If True, a learned head on the LSTM output produces the standard deviation. Otherwise the model uses an independent log-std parameter. Only used for continuous actions.

  • carry_leading_shape (tuple[int, ...]) – Leading dimensions for the persistent carry tensors h and c. Typically () at construction; the model resizes them lazily at runtime to match the batch shape.

  • discrete (bool) – If True, use a categorical distribution for discrete actions. If False (default), use a Gaussian distribution for continuous actions.

Variables:

  • obs_dim (int) – Observation dimensionality expected on the last axis of x.

  • lstm_features (int) – LSTM hidden/state size.

  • encoder (nnx.Sequential) – MLP that maps obs_dim → hidden_features.

  • cell (rnn.OptimizedLSTMCell) – LSTM cell with in_features = hidden_features and hidden_features = lstm_features.

  • fused_head (nnx.Linear) – Linear layer whose first action_space_size outputs are policy means (continuous) or logits (discrete), and whose last output is value.

  • actor_sigma (Callable[[jax.Array], jax.Array]) – Maps LSTM features to the policy standard deviations (learned head when actor_sigma_head=True, else independent parameter). Only used for continuous actions.

  • bij (distrax.Bijector) – Action-space bijector. Scalar bijectors are automatically lifted with Block(ndims=1) for vector actions (continuous only).

  • c (h,) – Persistent LSTM carry used by single-step evaluation. Shapes are (..., lstm_features). The model resizes them lazily to match the leading batch/agent dimensions.

  • discrete (bool) – Whether this model uses discrete actions.

property observation_space_size: int[source]#
property carry_leading_shape: tuple[int, ...][source]#
property carry: Any[source]#
reset(shape: tuple[int, ...], mask: Array | None = None) → None[source]#

Reset the persistent LSTM carry.

  • If self.h.value.shape != (*lead_shape, H), allocate fresh zeros once.

  • Otherwise, zero in-place:

    • if mask is None: zero everything

    • if mask is provided: zero selected environment/agent entries

Parameters:

  • shape (tuple[int, ...]) – Shape of the observation (input) tensor.

  • mask (optional bool array) – Boolean mask over leading environment/agent axes, broadcast over the remaining carry axes.

policy_with_perturbation(x: Array, perturbed_x: Array, **kwargs: Any) → tuple[Distribution, Array, Array][source]#

Evaluate spatial CAPS with the same clean incoming recurrent state.

See Model.policy_with_perturbation for the contract and the CAPS paper citation. Perturbations never advance the clean carry.