jaxdem.rl.environments#

Reinforcement-learning environment interface.

Classes

Environment(state, system, env_params)

Defines the interface for reinforcement-learning environments.

class jaxdem.rl.environments.Environment(state: State, system: System, env_params: dict[str, Any])#

Bases: Factory, ABC

Defines the interface for reinforcement-learning environments.

Let A be the number of agents (A ≥ 1). Single-agent environments still use A=1.
Observations and actions are flattened per agent to fixed sizes. Use action_space_shape to reshape inside the environment if needed.

Required shapes

Observation: (A, observation_space_size)
Action (input to step()): (A, action_space_size)
Reward: (A,)
Done: scalar boolean for the whole environment

Todo: - Truncated data field: per-agent termination flag - Render method

Example:#

To define a custom environment, inherit from Environment and implement the abstract methods:

>>> @Environment.register("MyCustomEnv")
>>> @jax.tree_util.register_dataclass
>>> @dataclass(slots=True)
>>> class MyCustomEnv(Environment):
    ...

state: State#: Simulation state.

system: System#: Simulation system configuration.

env_params: dict[str, Any]#: Environment-specific parameters.

classmethod Create(dim: int = 2) → Environment[source]#

Initialize the environment to a valid start state.

Parameters:

env ('MyCustomEnv') – Instance of the environment.
key (jax.random.PRNGKey) – JAX random number generator key.

Returns:

Freshly initialized environment.

Return type:

Environment

Conditionally resets the environment if the environment has reached a terminal state.

This method checks the done flag and, if True, calls the environment’s reset method to reinitialize the state. Otherwise, it returns the current environment unchanged.

Parameters:

env (Environment) – The current environment instance.
done (jax.Array) – A boolean flag indicating whether the environment has reached a terminal state.
key (jax.random.PRNGKey) – JAX random number generator key used for reinitialization.

Returns:

Either the freshly reset environment (if done is True) or the unchanged environment (if done is False).

Return type:

Environment

static step(env: Environment, action: Array) → Environment[source]#

Advance the simulation by one step using per-agent actions.

Parameters:

env (Environment) – The current environment.
action (jax.Array) – The vector of actions each agent in the environment should take.

Returns:

The updated environment state.

Return type:

Environment

static observation(env: Environment) → Array[source]#

Returns the per-agent observation vector.

Parameters:: env (Environment) – The current environment.
Returns:: Vector corresponding to the environment observation.
Return type:: jax.Array

static reward(env: Environment) → Array[source]#

Returns the per-agent immediate rewards.

Parameters:: env (Environment) – The current environment.
Returns:: Vector corresponding to all the agent’s rewards based on the current environment state.
Return type:: jax.Array

static done(env: Environment) → Array[source]#

Returns a boolean indicating whether the environment has ended.

Parameters:: env (Environment) – The current environment.
Returns:: A bool indicating when the environment ended
Return type:: jax.Array

static info(env: Environment) → dict[str, Any][source]#

Return auxiliary diagnostic information.

By default, returns an empty dict. Subclasses may override to provide environment specific information.

Parameters:: env (Environment) – The current state of the environment.
Returns:: A dictionary with additional information about the environment.
Return type:: Dict

property num_envs: int[source]#: Number of batched environments.

property max_num_agents: int[source]#: Maximum number of active agents in the environment.

property action_space_size: int[source]#: Flattened action size per agent. Actions passed to step() have shape (A, action_space_size).

property action_space_shape: tuple[int][source]#: Original per-agent action shape (useful for reshaping inside the environment).

property observation_space_size: int[source]#: Flattened observation size per agent. observation() returns shape (A, observation_space_size).

class jaxdem.rl.environments.MultiNavigator(state: State, system: System, env_params: dict[str, Any], n_lidar_rays: int)#

Bases: Environment

Multi-agent navigation environment toward assigned targets.

Each agent controls a force vector that is applied directly to a sphere inside a reflective box. Viscous drag -friction * vel is added each step. Objectives are sampled and assigned one-to-one via a random permutation.

The reward uses potential-based shaping with a proximity-gated kinetic-energy term:

\[\varphi_i(d, K) = \exp\!\left(-2 d^{\mathrm{eff}} - \frac{K}{\text{ke\_tau}}\,e^{-\text{ke\_gate} \cdot d^{\mathrm{eff}}}\right)\]

where \(d^{\mathrm{eff}} = \max(0, d - 0.5 r)\), \(d\) is the distance to the assigned objective, \(K\) is the translational kinetic energy, ke_tau sets the overall strength of the KE penalty, and ke_gate controls how sharply KE sensitivity falls off with distance — larger ke_gate means KE only matters very close to the objective. The per-agent shaping credit is \(F_i = \varphi_i(d^{\mathrm{eff}}_t, K_t) - \varphi_i(d^{\mathrm{eff}}_{t-1}, K_{t-1})\).

Notes

The observation vector per agent is:

Feature	Size
Unit direction to objective	`dim`
Clamped displacement	`dim`
Velocity	`dim`
LiDAR proximity (normalised)	`n_lidar_rays`

If one wants some realistic parameters for training, skip_frames = 50 will give a response rate of 200 Hz, meaning that num_steps_epoch = 100 gives a horizon of 0.5 seconds.

n_lidar_rays: int#: Number of angular bins for each LiDAR sensor.

classmethod Create(N: int = 64, min_box_size: float = 20.0, max_box_size: float = 20.0, box_padding: float = 5.0, max_steps: int = 100000, friction: float = 0.2, ke_tau: float = 5.0, ke_gate: float = 4.0, near_goal_bonus: float = 0.1, lidar_range: float = 10.0, n_lidar_rays: int = 16) → MultiNavigator[source]#

Create a multi-agent navigator environment.

Parameters:

N (int) – Number of agents.
min_box_size (float) – Range for the random square domain side length sampled at each reset().
max_box_size (float) – Range for the random square domain side length sampled at each reset().
box_padding (float) – Extra padding around the domain in multiples of the particle radius.
max_steps (int) – Episode length in physics steps.
friction (float) – Viscous drag coefficient applied as -friction * vel.
ke_tau (float) – Overall strength of the KE term in the potential (larger = less important). See class docstring.
ke_gate (float) – Distance decay rate of KE sensitivity (larger = KE only matters very close to the goal). See class docstring.
near_goal_bonus (float) – Reward bonus applied when an agent is within one radius of its objective.
lidar_range (float) – Maximum detection range for the LiDAR sensor.
n_lidar_rays (int) – Number of angular LiDAR bins spanning \([-\pi, \pi)\).

Returns:

A freshly constructed environment (call reset() before use).

Return type:

MultiNavigator

Initialize the environment with random positions and objectives.

Parameters:

env (Environment) – Current environment instance.
key (ArrayLike) – JAX random number generator key.

Returns:

Freshly initialized environment.

Return type:

Environment

static step(env: MultiNavigator, action: Array) → Environment[source]#

Advance one step. Actions are forces; simple drag is applied (-friction * vel).

Parameters:

env (Environment) – The current environment.
action (jax.Array) – The vector of actions each agent in the environment should take.

Returns:

The updated environment state.

Return type:

Environment

static observation(env: MultiNavigator) → Array[source]#

Build per-agent observations.

Contents per agent#

Unit vector to objective (shape (dim,)) –> Direction
Clamped delta to objective (shape (dim,)) –> Local precision
Velocity (shape (dim,))
LiDAR proximity, normalized by lidar_range (shape (n_lidar_rays,))

returns:: Array of shape (N, 3 * dim + n_lidar_rays)
rtype:: jax.Array

static reward(env: MultiNavigator) → Array[source]#

Returns a vector of per-agent rewards.

Potential-based shaping with a proximity-gated KE term:

\[\varphi(d, K) = \exp\!\left(-2 d^{\mathrm{eff}} - \frac{K}{\text{ke\_tau}}\,e^{-\text{ke\_gate} \cdot d^{\mathrm{eff}}}\right)\]

The gate \(e^{-\text{ke\_gate} \cdot d^{\mathrm{eff}}}\) suppresses the KE term away from the objective, so fast motion is free until the agent is close; ke_tau sets the overall strength of the penalty.

Per-step reward:

\[\mathrm{rew}_t = \frac{F_t + w_{\text{near}} \cdot \mathbf{1}[d_t \le r]}{w_{\text{near}}}\]

where \(F_t = \varphi(d^{\mathrm{eff}}_t, K_t) - \varphi(d^{\mathrm{eff}}_{t-1}, K_{t-1})\), \(d^{\mathrm{eff}}_t = \max(0, d_t - 0.5 r)\), and \(w_{\text{near}}\) weights a near-goal bonus.

Parameters:: env (Environment) – Current environment.
Returns:: Shape (N,).
Return type:: jax.Array

static done(env: MultiNavigator) → Array[source]#

Returns a boolean indicating whether the environment has ended. The episode terminates when the maximum number of steps is reached.

Parameters:: env (Environment) – The current environment.
Returns:: Boolean array indicating whether the episode has ended.
Return type:: jax.Array

property action_space_size: int[source]#: Flattened action size per agent. Actions passed to step() have shape (A, action_space_size).

property action_space_shape: tuple[int][source]#: Original per-agent action shape (useful for reshaping inside the environment).

property observation_space_size: int[source]#: Flattened observation size per agent. observation() returns shape (A, observation_space_size).

class jaxdem.rl.environments.MultiRoller(state: State, system: System, env_params: dict[str, Any], n_lidar_rays: int)#

Bases: Environment

Multi-agent rolling environment toward assigned targets.

Each agent controls a torque vector that is applied directly to a sphere on a \(z=0\) floor. Translational drag -friction * vel and angular damping -friction * ang_vel are applied each step. Objectives are sampled and assigned one-to-one via a random permutation.

The reward uses potential-based shaping with a proximity-gated kinetic-energy term:

\[\varphi(d, K) = \exp\!\left(-2 d^{\mathrm{eff}} - \frac{K}{\text{ke\_tau}}\,e^{-\text{ke\_gate} \cdot d^{\mathrm{eff}}}\right)\]

where \(d^{\mathrm{eff}} = \max(0, d - 0.5 r)\), \(d\) is the distance to the assigned objective in the \(xy\) plane, \(K\) is the translational kinetic energy, ke_tau sets the overall strength of the KE penalty, and ke_gate controls how sharply KE sensitivity falls off with distance — larger ke_gate means KE only matters very close to the objective. The per-agent shaping credit is \(F_i = \varphi(d^{\mathrm{eff}}_t, K_t) - \varphi(d^{\mathrm{eff}}_{t-1}, K_{t-1})\).

Notes

The observation vector per agent is:

Feature	Size
Unit direction to objective	`2`
Clamped displacement	`2`
Velocity	`2`
LiDAR proximity (normalised)	`n_lidar_rays`

If one wants some realistic parameters for training, skip_frames = 50 will give a response rate of 200 Hz, meaning that num_steps_epoch = 100 gives a horizon of 0.5 seconds.

n_lidar_rays: int#: Number of angular bins for each LiDAR sensor.

classmethod Create(N: int = 64, min_box_size: float = 20.0, max_box_size: float = 20.0, box_padding: float = 5.0, max_steps: int = 100000, friction: float = 0.2, ke_tau: float = 5.0, ke_gate: float = 4.0, near_goal_bonus: float = 0.1, lidar_range: float = 6.0, n_lidar_rays: int = 16) → MultiRoller[source]#

Create a multi-agent roller environment.

Parameters:

N (int) – Number of agents.
min_box_size (float) – Range for the random square domain side length sampled at each reset().
max_box_size (float) – Range for the random square domain side length sampled at each reset().
box_padding (float) – Extra padding around the domain in multiples of the particle radius.
max_steps (int) – Episode length in physics steps.
friction (float) – Translational and angular damping coefficient.
ke_tau (float) – Overall strength of the KE term in the potential (larger = less important). See class docstring.
ke_gate (float) – Distance decay rate of KE sensitivity (larger = KE only matters very close to the goal). See class docstring.
near_goal_bonus (float) – Reward bonus applied when an agent is within one radius of its objective.
lidar_range (float) – Maximum detection range for the LiDAR sensor.
n_lidar_rays (int) – Number of angular LiDAR bins spanning \([-\pi, \pi)\).

Returns:

A freshly constructed environment (call reset() before use).

Return type:

MultiRoller

Initialize the environment with random positions and objectives.

Parameters:

env (Environment) – Current environment instance.
key (ArrayLike) – JAX random number generator key.

Returns:

Freshly initialized environment.

Return type:

Environment

static step(env: MultiRoller, action: Array) → Environment[source]#

Advance one step. Actions are torques; simple damping is applied.

Parameters:

env (Environment) – The current environment.
action (jax.Array) – The vector of actions each agent in the environment should take.

Returns:

The updated environment state.

Return type:

Environment

static observation(env: MultiRoller) → Array[source]#

Build per-agent observations.

Contents per agent#

Unit vector to objective in the \(xy\) plane (shape (2,)).
Clamped objective delta in the \(xy\) plane (shape (2,)).
Velocity in the \(xy\) plane (shape (2,)).
LiDAR proximity, normalized by lidar_range (shape (n_lidar_rays,)).

returns:: Array of shape (N, 6 + n_lidar_rays)
rtype:: jax.Array

static reward(env: MultiRoller) → Array[source]#

Returns a vector of per-agent rewards.

Potential-based shaping with a proximity-gated KE term:

\[\varphi(d, K) = \exp\!\left(-2 d^{\mathrm{eff}} - \frac{K}{\text{ke\_tau}}\,e^{-\text{ke\_gate} \cdot d^{\mathrm{eff}}}\right)\]

Per-step reward:

\[\mathrm{rew}_t = \frac{F_t + w_{\text{near}} \cdot \mathbf{1}[d_t \le r]}{w_{\text{near}}}\]

where \(F_t = \varphi(d^{\mathrm{eff}}_t, K_t) - \varphi(d^{\mathrm{eff}}_{t-1}, K_{t-1})\), \(d^{\mathrm{eff}}_t = \max(0, d_t - 0.5 r)\), and \(w_{\text{near}}\) weights a near-goal bonus.

Parameters:: env (Environment) – Current environment.
Returns:: Shape (N,).
Return type:: jax.Array

static done(env: MultiRoller) → Array[source]#

Returns a boolean indicating whether the environment has ended. The episode terminates when the maximum number of steps is reached.

Parameters:: env (Environment) – The current environment.
Returns:: Boolean array indicating whether the environment has ended.
Return type:: jax.Array

property action_space_size: int[source]#: Flattened action size per agent. Actions passed to step() have shape (A, action_space_size).

property action_space_shape: tuple[int][source]#: Original per-agent action shape (useful for reshaping inside the environment).

property observation_space_size: int[source]#: Flattened observation size per agent. observation() returns shape (A, observation_space_size).

class jaxdem.rl.environments.SingleNavigator(state: State, system: System, env_params: dict[str, Any])#

Bases: Environment

Single-agent navigation environment toward a fixed target.

The agent controls a force vector that is applied directly to a sphere inside a reflective box. Viscous drag -friction * vel is added each step. The reward uses potential-based shaping with a proximity-gated kinetic-energy term:

\[\varphi(d, K) = \exp\!\left(-2 d - \frac{K}{\text{ke\_tau}}\,e^{-\text{ke\_gate} \cdot d}\right)\]

where \(d\) is the distance to the objective, \(K\) is the translational kinetic energy, ke_tau is the KE scale that sets the overall strength of the penalty, and ke_gate controls how sharply KE sensitivity falls off with distance — larger ke_gate means KE only matters very close to the objective.

The shaping credit is \(F_t = \varphi(d_t, K_t) - \varphi(d_{t-1}, K_{t-1})\), so kinetic energy is penalised only near the objective — far away the gate \(e^{-\text{ke\_gate} \cdot d} \to 0\) and fast motion is free.

Per-step reward:

\[\mathrm{rew}_t = \frac{F_t + b \cdot \mathbb{1}[d_t \le r]}{b}\]

where \(b\) is the near-goal bonus and \(r\) is the agent radius.

Notes

The observation vector per agent is:

Feature	Size
Unit direction to objective	`dim`
Clamped displacement	`dim`
Velocity	`dim`

If one wants some realistic parameters for training, skip_frames = 50 will give a response rate of 200 Hz, meaning that num_steps_epoch = 100 gives a horizon of 0.5 seconds.

classmethod Create(dim: int = 2, min_box_size: float = 40.0, max_box_size: float = 40.0, max_steps: int = 20000, friction: float = 0.2, near_goal_bonus: float = 0.1, ke_tau: float = 2.0, ke_gate: float = 6.0) → SingleNavigator[source]#

Create a single-agent navigator environment.

Parameters:

dim (int) – Spatial dimensionality (2 or 3).
min_box_size (float) – Range for the random square domain side length.
max_box_size (float) – Range for the random square domain side length.
max_steps (int) – Episode length in physics steps.
friction (float) – Viscous drag coefficient applied as -friction * vel.
ke_tau (float) – Overall strength of the KE term in the potential (larger = less important). See class docstring.
ke_gate (float) – Distance decay rate of KE sensitivity (larger = KE only matters very close to the goal). See class docstring.

Returns:

A freshly constructed environment (call reset() before use).

Return type:

SingleNavigator

Initialize the environment with a randomly placed particle and velocity.

Parameters:

env ('SingleNavigator') – Current environment instance.
key (jax.random.PRNGKey) – JAX random number generator key.

Returns:

Freshly initialized environment.

Return type:

Environment

static step(env: SingleNavigator, action: Array) → Environment[source]#

Advance one step. Actions are forces; simple drag is applied (-friction * vel).

Parameters:

env (Environment) – The current environment.
action (jax.Array) – The vector of actions each agent in the environment should take.

Returns:

The updated environment state.

Return type:

Environment

static observation(env: SingleNavigator) → Array[source]#

Build per-agent observations.

Contents per agent#

Unit vector to objective (shape (dim,)) –> Direction
Clamped delta to objective (shape (dim,)) –> Local precision
Velocity (shape (dim,))

returns:: Array of shape (N, 3 * dim)
rtype:: jax.Array

static reward(env: SingleNavigator) → Array[source]#

Returns a vector of per-agent rewards.

Potential-based shaping with a proximity-gated KE term:

\[\varphi(d, K) = \exp\!\left(-2 d - \frac{K}{\text{ke\_tau}}\,e^{-\text{ke\_gate} \cdot d}\right)\]

The gate \(e^{-\text{ke\_gate} \cdot d}\) suppresses the KE term away from the objective, so fast motion is free until the agent is close; ke_tau sets the overall strength of the penalty.

Per-step reward:

\[\mathrm{rew}_t = \frac{\varphi(d_t, K_t) - \varphi(d_{t-1}, K_{t-1}) + b \cdot \mathbb{1}[d_t \le r]}{b}\]

where \(b\) is the near-goal bonus and \(r\) is the agent radius.

Parameters:: env (Environment) – Current environment.
Returns:: Shape (N,).
Return type:: jax.Array

static done(env: SingleNavigator) → Array[source]#

Returns a boolean indicating whether the environment has ended. The episode terminates when the maximum number of steps is reached.

Parameters:: env (Environment) – The current environment.
Returns:: Boolean array indicating whether the episode has ended.
Return type:: jax.Array

property action_space_size: int[source]#: Flattened action size per agent. Actions passed to step() have shape (A, action_space_size).

property action_space_shape: tuple[int][source]#: Original per-agent action shape (useful for reshaping inside the environment).

property observation_space_size: int[source]#: Flattened observation size per agent. observation() returns shape (A, observation_space_size).

class jaxdem.rl.environments.SingleRoller(state: State, system: System, env_params: dict[str, Any])#

Bases: Environment

Single-agent 3D navigation via torque-controlled rolling.

The agent is a sphere resting on a \(z = 0\) floor under gravity. Actions are 3-D torque vectors; translational motion arises from frictional contact with the floor (see frictional_wall_force()). A viscous drag -friction * vel and a fixed angular damping of -friction * ang_vel are applied each step.

The reward uses potential-based shaping with a proximity-gated kinetic-energy term:

\[\varphi(d, K) = \exp\!\left(-2 d - \frac{K}{\text{ke\_tau}}\,e^{-\text{ke\_gate} \cdot d}\right)\]

where \(d\) is the distance to the objective, \(K\) is the total (translational + rotational) kinetic energy, ke_tau is the KE scale that sets the overall strength of the penalty, and ke_gate controls how sharply KE sensitivity falls off with distance — larger ke_gate means KE only matters very close to the objective.

Per-step reward:

\[\mathrm{rew}_t = \frac{F_t + b \cdot \mathbb{1}[d_t \le r]}{b}\]

where \(b\) is the near-goal bonus and \(r\) is the agent radius.

Notes

The observation vector per agent is:

Feature	Size
Unit direction to objective	2
Clamped displacement (x, y)	2
Velocity (x, y)	2
Angular velocity	3

If one wants some realistic parameters for training, skip_frames = 50 will give a response rate of 200 Hz, meaning that num_steps_epoch = 100 gives a horizon of 0.5 seconds.

classmethod Create(min_box_size: float = 40.0, max_box_size: float = 40.0, max_steps: int = 20000, friction: float = 0.2, near_goal_bonus: float = 0.1, ke_tau: float = 5.0, ke_gate: float = 4.0) → SingleRoller[source]#

Create a single-agent roller environment.

Parameters:

min_box_size (float) – Range for the random square domain side length.
max_box_size (float) – Range for the random square domain side length.
max_steps (int) – Episode length in physics steps.
friction (float) – Viscous drag coefficient applied as -friction * vel.
ke_tau (float) – Overall strength of the KE term in the potential (larger = less important). See class docstring.
ke_gate (float) – Distance decay rate of KE sensitivity (larger = KE only matters very close to the goal). See class docstring.

Returns:

A freshly constructed environment (call reset() before use).

Return type:

SingleRoller

Randomly place the agent and objective on the floor.

Parameters:

env (Environment) – Current environment instance.
key (ArrayLike) – JAX PRNG key.

Returns:

Freshly initialised environment.

Return type:

Environment

static step(env: SingleRoller, action: Array) → Environment[source]#

Apply a torque action, advance physics by one step.

Parameters:

env (Environment) – Current environment.
action (jax.Array) – 3-D torque vector per agent.

Returns:

Updated environment after one physics step.

Return type:

Environment

static observation(env: SingleRoller) → Array[source]#

Per-agent observation vector.

Contents per agent:

Unit displacement to objective projected to x-y (shape (2,)).
Clamped displacement to objective projected to x-y (shape (2,)).
Velocity projected to x-y (shape (2,)).
Angular velocity (shape (3,)).

Returns:: Shape (N, 9).
Return type:: jax.Array

static reward(env: SingleRoller) → Array[source]#

Returns a vector of per-agent rewards.

Potential-based shaping with a proximity-gated KE term:

\[\varphi(d, K) = \exp\!\left(-2 d - \frac{K}{\text{ke\_tau}}\,e^{-\text{ke\_gate} \cdot d}\right)\]

The gate \(e^{-\text{ke\_gate} \cdot d}\) suppresses the KE term away from the objective, so fast motion is free until the agent is close; ke_tau sets the overall strength of the penalty.

Per-step reward:

\[\mathrm{rew}_t = \frac{\varphi(d_t, K_t) - \varphi(d_{t-1}, K_{t-1}) + b \cdot \mathbb{1}[d_t \le r]}{b}\]

where \(b\) is the near-goal bonus and \(r\) is the agent radius.

Returns:: Shape (N,).
Return type:: jax.Array

static done(env: SingleRoller) → Array[source]#: True when step_count exceeds max_steps.

property action_space_size: int[source]#: Per-agent flattened action dimensionality (3-D torque).

property action_space_shape: tuple[int][source]#: Per-agent action tensor shape.

property observation_space_size: int[source]#: Per-agent flattened observation dimensionality (9).

class jaxdem.rl.environments.SwarmNavigator(state: State, system: System, env_params: dict[str, Any], n_lidar_rays: int, num_objectives: int)#

Bases: Environment

Multi-agent cooperative objective coverage with local sensing.

Each agent controls a force vector applied to a sphere in a reflective box, with viscous drag -friction * vel added each step. Objectives are sampled on a jittered grid inside the box; agents spawn in the padding ring around it. Three LiDAR sensors are refreshed each step — walls, objectives, and peers (other agents) — but only the objective and wall sensors appear in the observation; the peer sensor drives the contention penalty in the reward. lidar_obj_prev and lidar_agt_prev hold the previous step’s objective and peer readings so the reward can difference them.

Notes

The observation vector per agent is:

Feature	Size
Velocity	`dim`
Objective LiDAR (normalised)	`n_lidar_rays`
Wall LiDAR (normalised)	`n_lidar_rays`

n_lidar_rays: int#: Number of angular bins for each LiDAR sensor.

num_objectives: int#: Number of objectives sampled per environment.

classmethod Create(N: int = 64, num_objectives: int = 64, box_size: float = 20.0, box_padding: float = 10.0, max_steps: int = 10000, friction: float = 0.2, near_goal_bonus: float = 0.01, lidar_range: float = 16.0, n_lidar_rays: int = 12, contention_strength: float = 15.0) → SwarmNavigator[source]#

Create a swarm navigator environment.

Parameters:

N (int) – Number of agents.
num_objectives (int) – Number of objectives sampled per environment.
box_size (float) – Side length of the square domain that holds the objectives.
box_padding (float) – Thickness of the agent spawn ring around the box (in multiples of the particle radius).
max_steps (int) – Episode length in physics steps.
friction (float) – Viscous drag coefficient applied as -friction * vel.
near_goal_bonus (float) – Weight \(b\) of the near-goal indicator \(\mathbf{1}[d \le r]\).
lidar_range (float) – Maximum detection range \(L\) for the LiDAR sensors.
n_lidar_rays (int) – Number of angular LiDAR bins spanning \([-\pi, \pi)\).
contention_strength (float) – Maximum penalty \(P_{\max}\) subtracted from an objective’s LiDAR proximity when a peer sits on it; the bin-wise penalty ramps linearly from \(P_{\max}\) (peer on the objective) to 0 (peer at \(L/4\)) and is zero beyond.

Returns:

A freshly constructed environment (call reset() before use).

Return type:

SwarmNavigator

static reset(env: SwarmNavigator, key: Array | ndarray | bool | number | bool | int | float | complex) → Environment[source]#: Initialise the environment with random agents (padding) and objectives (box).

static step(env: SwarmNavigator, action: Array) → Environment[source]#: Advance one step. Actions are forces; drag -friction * vel is added.

static observation(env: SwarmNavigator) → Array[source]#: Velocity + objective LiDAR + wall LiDAR (all normalised), per agent.

static reward(env: SwarmNavigator) → Array[source]#

Potential-based shaping with a bin-wise contention penalty.

For each objective LiDAR bin, the nearest agent (over all agent LiDAR bins, distance recovered with the law of cosines) subtracts from the objective’s apparent proximity when it lies within lr/4 of it (exponential decay, already negligible by lr/4):

d_eff = d_obj + P_max * exp(-d_peer / tau),  tau = 1.0

where d_peer is \(\min_a \sqrt{d_{obj}^2 + d_{agt,a}^2 - 2 d_{obj} d_{agt,a} \cos(\Delta\theta)}`\). Empty bins read at lr (max range); the resulting long-range inaccuracy is negligible since far objectives barely contribute.

Per-step reward:

R = near_goal_bonus * 1[d_min <= r] + 10 * (phi_t - phi_prev)

where d_min is the closest objective distance and 10 is the shaping scale.

static done(env: SwarmNavigator) → Array[source]#: Episode terminates when max_steps is reached.

property action_space_size: int[source]#: Flattened action size per agent.

property action_space_shape: tuple[int][source]#: Original per-agent action shape.

property observation_space_size: int[source]#: Flattened observation size per agent.

class jaxdem.rl.environments.SwarmRoller(state: State, system: System, env_params: dict[str, Any], n_lidar_rays: int, num_objectives: int)#

Bases: Environment

Multi-agent cooperative objective coverage with rolling dynamics.

Each agent controls a torque vector applied to a sphere on a \(z=0\) floor, with translational drag -friction * vel and angular damping -friction * ang_vel added each step. Objectives are sampled on a jittered grid inside the box (at floor level); agents spawn in the padding ring around it. Three LiDAR sensors are refreshed each step — walls, objectives, and peers (other agents) — but only the objective and wall sensors appear in the observation; the peer sensor drives the contention penalty in the reward. lidar_obj_prev and lidar_agt_prev hold the previous step’s objective and peer readings so the reward can difference them.

Notes

The observation vector per agent is:

Feature	Size
Velocity	`dim`
Angular velocity	`dim`
Objective LiDAR (normalised)	`n_lidar_rays`
Wall LiDAR (normalised)	`n_lidar_rays`

n_lidar_rays: int#: Number of angular bins for each LiDAR sensor.

num_objectives: int#: Number of objectives sampled per environment.

Create a swarm roller environment.

Parameters:

N (int) – Number of agents.
num_objectives (int) – Number of objectives sampled per environment.
box_size (float) – Side length of the square domain that holds the objectives.
box_padding (float) – Thickness of the agent spawn ring around the box (in multiples of the particle radius).
max_steps (int) – Episode length in physics steps.
friction (float) – Translational and angular damping applied as -friction * vel and -friction * ang_vel.
near_goal_bonus (float) – Weight \(b\) of the near-goal indicator \(\mathbf{1}[d \le r]\).
lidar_range (float) – Maximum detection range \(L\) for the LiDAR sensors.
n_lidar_rays (int) – Number of angular LiDAR bins spanning \([-\pi, \pi)\).
contention_strength (float) – Maximum penalty \(P_{\max}\) subtracted from an objective’s LiDAR proximity when a peer sits on it; the bin-wise penalty ramps linearly from \(P_{\max}\) (peer on the objective) to 0 (peer at \(L/4\)) and is zero beyond.

Returns:

A freshly constructed environment (call reset() before use).

Return type:

SwarmRoller

static reset(env: SwarmRoller, key: Array | ndarray | bool | number | bool | int | float | complex) → Environment[source]#: Initialise the environment with random agents (padding) and objectives (box).

static step(env: SwarmRoller, action: Array) → Environment[source]#: Advance one step. Actions are torques; drag -friction * vel and -friction * ang_vel are added.

static observation(env: SwarmRoller) → Array[source]#: Velocity + angular velocity + objective LiDAR + wall LiDAR (all normalised), per agent.

static reward(env: SwarmRoller) → Array[source]#

Potential-based shaping with a bin-wise contention penalty.

d_eff = d_obj + P_max * exp(-d_peer / tau),  tau = 1.0

Per-step reward:

R = near_goal_bonus * 1[d_min <= r] + 10 * (phi_t - phi_prev)

where d_min is the closest objective distance and 10 is the shaping scale.

static done(env: SwarmRoller) → Array[source]#: Episode terminates when max_steps is reached.

property action_space_size: int[source]#: Flattened action size per agent (torque components).

property action_space_shape: tuple[int][source]#: Original per-agent action shape.

property observation_space_size: int[source]#: Flattened observation size per agent.

class jaxdem.rl.environments.SwarmRoller3D(state: State, system: System, env_params: dict[str, Any], n_lidar_rays: int, n_lidar_elevation: int, num_objectives: int)#

Bases: Environment

Multi-agent cooperative coverage of 3-D pyramid objectives with attraction.

Identical in structure to SwarmRoller: rolling-sphere agents with translational and angular drag, three LiDAR sensors (walls, objectives, peers) and a bin-wise contention-shaped reward. Two differences: objectives are arranged as a square pyramid (sensed with 3-D LiDAR), and agents exert pairwise magnetic attraction on each other.

Feature	Size
Velocity	`dim`
Angular velocity	`dim`
Objective LiDAR (normalised)	`n_az * n_el`
Wall LiDAR (normalised)	`n_az * n_el`

n_lidar_rays: int#: Number of azimuthal bins for each 3-D LiDAR sensor.

n_lidar_elevation: int#: Number of elevation bins for each 3-D LiDAR sensor.

num_objectives: int#: Number of objectives (pyramid spheres) sampled per environment.

classmethod Create(N: int = 5, num_objectives: int = 5, box_size: float = 5.0, box_padding: float = 5.0, max_steps: int = 10000, friction: float = 0.2, near_goal_bonus: float = 0.01, lidar_range: float = 16.0, n_lidar_rays: int = 8, n_lidar_elevation: int = 8, contention_strength: float = 15.0, magnet_strength: float = 4.0, magnet_range: float = 3.0) → SwarmRoller3D[source]#

Create a 3-D swarm roller environment with pyramid objectives.

Parameters mirror SwarmRoller.Create(), plus n_lidar_elevation (3-D LiDAR elevation bins) and magnet_strength / magnet_range for the inter-agent attraction.

static reset(env: SwarmRoller3D, key: Array | ndarray | bool | number | bool | int | float | complex) → Environment[source]#: Initialise with agents in the padding ring and a pyramid of objectives in the box.

static step(env: SwarmRoller3D, action: Array) → Environment[source]#: Advance one step: drag, torque, mutual attraction, then physics + sensing.

static observation(env: SwarmRoller3D) → Array[source]#: Velocity + angular velocity + objective LiDAR + wall LiDAR (normalised), per agent.

static reward(env: SwarmRoller3D) → Array[source]#

Potential-based shaping with a bin-wise contention penalty.

Same as SwarmRoller.reward(), but the law-of-cosines bin geometry uses azimuth alignment (az = bin // n_elevation) since the bins are the flattened 3-D (azimuth, elevation) grid.

static done(env: SwarmRoller3D) → Array[source]#: Episode terminates when max_steps is reached.

property action_space_size: int[source]#: Flattened action size per agent (torque components).

property action_space_shape: tuple[int][source]#: Original per-agent action shape.

property observation_space_size: int[source]#: Flattened observation size per agent.

class jaxdem.rl.environments.ThreeGears(state: State, system: System, env_params: dict[str, Any], num_gears: int)#

Bases: Environment

N dynamic gears that must assemble a triangular stack.

Identical dynamics, pairwise attraction, nearest-neighbour observation, and per-gear reward as TwoGears — only the objective differs: the num_gears targets form a triangular stack (rows that shrink by one from bottom to top, gears touching). num_gears=3 is the classic triangle [2,1]; 5 -> [3,2]; 6 -> [3,2,1]. Gear i is paired with objective i.

Note

As with TwoGears, skip_frames = 50 gives a 200 Hz response rate, so num_steps_epoch = 100 is a 0.5 s horizon. box_size must fit the stack — width 2*m*rr and height (2 + (m-1)*sqrt(3))*rr (m = bottom row size), and >= 2*rr*(num_gears+1) wide for a non-overlapping spawn.

num_gears: int#: Number of gears (agents) forming the triangular stack.

classmethod Create(num_gears: int = 6, box_size: float = 30.0, max_steps: int = 100000, friction: float = 0.2, ke_weight: float = 0.1, attraction_mag: float = 2.0) → ThreeGears[source]#

Create an N-gear triangular-stack environment.

Parameters:

num_gears (int) – Number of dynamic gears (agents) forming the stack.
box_size (float) – Size of the square bounding box.
max_steps (int) – Episode length in physics steps.
friction (float) – Viscous drag coefficient applied as -friction * vel.
ke_weight (float) – Weight for the differential kinetic energy penalty.
attraction_mag (float) – Magnitude of the pairwise attraction force between gears.

Returns:

A freshly constructed environment (call reset() before use).

Return type:

ThreeGears

static reset(env: ThreeGears, key: Array) → Environment[source]#: Reset with the gears on the floor and a random triangular-stack objective.

static step(env: ThreeGears, action: Array) → Environment[source]#

Advance one step: per-gear torque, pairwise attraction, viscous drag.

Attraction on gear \(i\) from gear \(j\) is \(-(C/d_{ij}^3)\,\hat{n}_{ij}\) when \(d_{ij} < 3r\), with \(\hat{n}_{ij}=\mathrm{unit}(\mathbf{r}_i-\mathbf{r}_j)\) and \(C = m_{\text{attr}}(2r)^3\). Net force on \(i\) is \(\sum_{j\ne i}\).

static observation(env: ThreeGears) → Array[source]#

Per-gear observation (16 features); “other gear” = nearest neighbour.

Feature	Size
Distance to floor	`1`
Distance to left/right walls	`2`
Unit vector to target	`2`
Clamped displacement to target	`2`
Unit vector to nearest gear	`2`
Clamped displacement to nearest gear	`2`
\(\sin(\Delta\theta)\)	`1`
\(\cos(\Delta\theta)\)	`1`
Velocity (x, y)	`2`
Angular velocity	`1`

static reward(env: ThreeGears) → Array[source]#: Per-gear shaping reward.

\[R_i = (d_{i,t-1} - d_{i,t}) - w_{\text{ke}} (K_{i,t} - K_{i,t-1})\]

static done(env: ThreeGears) → Array[source]#

property action_space_size: int[source]#: Flattened action size per agent. Actions passed to step() have shape (A, action_space_size).

property action_space_shape: tuple[int][source]#: Original per-agent action shape (useful for reshaping inside the environment).

property observation_space_size: int[source]#: Flattened observation size per agent. observation() returns shape (A, observation_space_size).

property max_num_agents: int[source]#: Maximum number of active agents in the environment.

class jaxdem.rl.environments.TwoGears(state: State, system: System, env_params: dict[str, Any], num_gears: int)#

Bases: Environment

Two-dimensional environment with N dynamic gears building a tower.

All num_gears gears are dynamic agents that each apply torque to themselves. Each episode samples a random target x and stacks num_gears objectives vertically into a tower (gear i must reach level i, bottom to top). The gears spawn at random, non-overlapping floor positions — not necessarily under the tower — and must navigate to assemble the stack. Gears attract each other pairwise via a magnetic force, and each gear observes its nearest neighbour.

Note

After experimentation, one needs the max torque to be at least 4.0 * mgr for the gear to be able to climb correctly, and attraction at least 1 * mg. If one wants some realistic parameters for training, skip_frames = 50 will give a response rate of 200 Hz, meaning that num_steps_epoch = 100 gives a horizon of 0.5 seconds. box_size must fit num_gears gears of radius rr side by side on the floor (box_size >= 2*rr*(num_gears+1)) and fit the tower height 2*rr*num_gears vertically.

num_gears: int#: Number of gears (agents) that must form the tower.

classmethod Create(num_gears: int = 3, box_size: float = 20.0, max_steps: int = 100000, friction: float = 0.2, ke_weight: float = 0.1, attraction_mag: float = 4.0) → TwoGears[source]#

Create an N-gear tower environment.

Parameters:

num_gears (int) – Number of dynamic gears (agents) that must form the tower.
box_size (float) – Size of the square bounding box.
max_steps (int) – Episode length in physics steps.
friction (float) – Viscous drag coefficient applied as -friction * vel.
ke_weight (float) – Weight for the differential kinetic energy penalty.
attraction_mag (float) – Magnitude of the pairwise attraction force between gears.

Returns:

A freshly constructed environment (call reset() before use).

Return type:

TwoGears

static reset(env: TwoGears, key: Array) → Environment[source]#

Reset the environment to a random initial configuration.

Parameters:

env (Environment) – The environment instance to reset.
key (jax.Array) – PRNG key used to sample the initial positions and objective.

Returns:

The environment with a fresh episode state.

Return type:

Environment

static step(env: TwoGears, action: Array) → Environment[source]#

Advance the environment by one step.

Applies each gear’s torque, computes the pairwise attraction force between all gears, and applies viscous drag.

The attraction on gear \(i\) from gear \(j\) is:

\[\mathbf{F}_{ij} = - \frac{C}{d_{ij}^3} \hat{n}_{ij},\]

when \(d_{ij} < 3 r\), where \(d_{ij}\) is the center-to-center distance, \(\hat{n}_{ij} = \mathrm{unit}(\mathbf{r}_i - \mathbf{r}_j)\) (so the force points from \(i\) toward \(j\)), and \(C = m_{\text{attr}} (2r)^3\) with \(r\) the gear radius. The net force on gear \(i\) is \(\sum_{j \ne i} \mathbf{F}_{ij}\).

Parameters:

env (Environment) – Current environment.
action (jax.Array) – Torque action for each gear, shape (num_gears, 1).

Returns:

Updated environment after physics integration and sensor updates.

Return type:

Environment

static observation(env: TwoGears) → Array[source]#

Build the per-gear observation vector.

Each gear receives a 16-feature observation; the “other gear” slot is filled by its nearest neighbour:

Feature	Size
Distance to floor	`1`
Distance to left/right walls	`2`
Unit vector to target	`2`
Clamped displacement to target	`2`
Unit vector to nearest gear	`2`
Clamped displacement to nearest gear	`2`
\(\sin(\Delta\theta)\)	`1`
\(\cos(\Delta\theta)\)	`1`
Velocity (x, y)	`2`
Angular velocity	`1`

Returns:: Observation of shape (num_gears, 16) — one row per gear.
Return type:: jax.Array

static reward(env: TwoGears) → Array[source]#

Compute the reward.

The reward is based on the differential distance to the objective minus a penalty for the change in kinetic energy:

\[R_t = (d_{t-1} - d_t) - w_{\text{ke}} (K_t - K_{t-1})\]

where \(d_t\) is the distance from gear \(i\) to its objective at step \(t\), \(K_t\) is that gear’s kinetic energy at step \(t\), and \(w_{\text{ke}}\) is the weight for the kinetic energy penalty.

Returns:: Per-gear reward of shape (num_gears,).
Return type:: jax.Array

static done(env: TwoGears) → Array[source]#

property action_space_size: int[source]#: Flattened action size per agent. Actions passed to step() have shape (A, action_space_size).

property action_space_shape: tuple[int][source]#: Original per-agent action shape (useful for reshaping inside the environment).

property observation_space_size: int[source]#: Flattened observation size per agent. observation() returns shape (A, observation_space_size).

property max_num_agents: int[source]#: Maximum number of active agents in the environment.

Modules

`burrowing_dp`
`hardcoded_bed`
`multi_navigator`	Environment where multiple agents navigate towards assigned targets.
`multi_roller`	Environment where multiple rolling agents navigate towards assigned targets.
`single_navigator`	Environment where a single agent navigates towards a target.
`single_roller`	Environment where a single agent rolls towards a target on the floor.
`swarm_navigator`	Environment where multiple agents cooperatively cover a set of objectives.
`swarm_roller`	Environment where multiple rolling agents cooperatively cover a set of objectives.
`swarm_roller_3d`	3-D swarm rolling agents covering pyramid objectives, with mutual attraction.
`three_gears`	Three-gear environment: three dynamic gears must assemble a triangle.
`two_gears`	Two-dimensional environment with two gears for RL training.