jaxdem.rl.models.mingru#

Reinforcement learning model based on MinGRU.

Classes

MinGRUActorCritic(*args, **kwargs)

A recurrent actor-critic with an MLP encoder and a MinGRU torso.
class jaxdem.rl.models.mingru.MinGRUActorCritic(*args: Any, **kwargs: Any)[source]#

Bases: Model

A recurrent actor-critic with an MLP encoder and a MinGRU torso.

This model uses the MinGRU architecture (https://arxiv.org/abs/2410.01201) and a parallel associative scan for sequence-mode training. Sequence mode uses initial_carry when provided, otherwise zeros. done[t] cuts the carry connection into step t+1; the mask is shifted internally. Sequence evaluation leaves the persistent rollout carry unchanged.

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.

  • gru_features (int) – MinGRU hidden size. Also the feature size consumed by the policy and value heads.

  • num_layers (int) – Number of MinGRU layers.

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

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

  • remat (bool) – Accepted for interface compatibility. Unused.

  • actor_sigma_head (bool) – If True, a learned head on the MinGRU 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 tensor h.

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

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

Reset the persistent recurrent carry.

Parameters:

  • shape (tuple[int, ...]) – Observation-shaped tuple, e.g. (num_envs, num_agents, obs_dim). Recurrent models use shape[:-1] as the carry leading shape; the final feature dimension may be a placeholder.

  • mask (optional bool array) – True where to reset entries, e.g. (num_envs,) or (num_envs, num_agents). Trailing singleton axes are added to broadcast over carry features. None resets all entries.

property carry: Any[source]#
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.