jaxdem.rl.models

Interface for defining reinforcement learning models.

Classes

Model(*args, **kwargs)	Base interface for reinforcement learning models.
SigmaHead(*args, **kwargs)	Predict positive action standard deviations from learned features.
SigmaParam(*args, **kwargs)	Learn a state-independent positive standard deviation for each action.

class jaxdem.rl.models.SigmaParam(*args: Any, **kwargs: Any)

Bases: Module

Learn a state-independent positive standard deviation for each action.

class jaxdem.rl.models.SigmaHead(*args: Any, **kwargs: Any)

Bases: Module

Predict positive action standard deviations from learned features.

class jaxdem.rl.models.Model(*args: Any, **kwargs: Any)

Bases: Factory, Module, ABC

Base interface for reinforcement learning models. Acts as a namespace.

Models map observations to an action distribution and a value estimate.

Example

To define a custom model, inherit from Model and implement its abstract methods:

>>> @Model.register("myCustomModel")
... class MyCustomModel(Model):
...     ...

discrete: bool

Whether the model emits a categorical (discrete) action distribution. Subclasses set this in __init__ before calling _init_policy_head().

actor_sigma: SigmaHead | SigmaParam

Continuous-policy scale, either a feature-dependent head or a learned state-independent parameter. Both accept the same feature tensor.

property action_space: Any

property carry: Any

static policy_output(pi: Distribution) → Array

Return differentiable policy outputs for smoothness regularization.

Continuous policies use the base mean passed through the action bijector, rather than the mean of the transformed distribution. Categorical policies use probabilities, avoiding an ordering of category labels and the zero gradients of argmax actions.

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

Evaluate a clean sequence and locally perturbed policy outputs.

Returns the clean policy, clean values, and perturbed outputs as defined by policy_output(). Recurrent implementations must compare both observations with the same clean incoming carry at each timestep, without changing persistent rollout state. This default implementation supports feedforward models.

Used for spatial CAPS: Mysore et al., Regularizing Action Policies for Smooth Control with Reinforcement Learning, ICRA 2021, https://arxiv.org/abs/2012.06644.

reset(shape: tuple[int, ...], mask: Array | None = None) → None

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.

Modules

lstm	Reinforcement learning model based on a single-layer LSTM.
mingru	Reinforcement learning model based on MinGRU.
mlp	Actor-critic reinforcement learning models based on MLPs.