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.
The environment flattens observations and actions 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. 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') – The current environment.
key (jax.random.PRNGKey) – JAX random number generator key.

Returns:

The initialized environment.

Return type:

Environment

Reset the environment when done is True.

When done is True, this method calls the environment’s reset method. When done is False, it returns the current environment unchanged.

Parameters:

env (Environment) – The current environment.
done (jax.Array) – A boolean flag that is True when the environment has reached a terminal state.
key (jax.random.PRNGKey) – JAX random number generator key used for the reset.

Returns:

The reset environment when done is True, otherwise the unchanged environment.

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 per-agent action vectors.

Returns:

The updated environment state.

Return type:

Environment

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

Return the per-agent observation vector.

Parameters:: env (Environment) – The current environment.
Returns:: Per-agent observations, shape (A, observation_space_size).
Return type:: jax.Array

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

Return the per-agent immediate rewards.

Parameters:: env (Environment) – The current environment.
Returns:: Per-agent rewards for the current state, shape (A,).
Return type:: jax.Array

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

Return whether the episode has ended.

Parameters:: env (Environment) – The current environment.
Returns:: A bool that is True when the episode has ended.
Return type:: jax.Array

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

Return auxiliary diagnostic information.

The default is an empty dict. Subclasses can override this method to provide environment-specific information.

Parameters:: env (Environment) – The current 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 acts directly on a sphere inside a reflective box. Each step adds viscous drag -friction * vel. The environment samples objectives and assigns them one-to-one with 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, and \(K\) is the translational kinetic energy. ke_tau sets the overall strength of the KE penalty. ke_gate controls how sharply KE sensitivity falls off with distance. A 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 (normalized)	`n_lidar_rays`

For realistic training parameters, skip_frames = 50 gives a response rate of 200 Hz, so 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:

The constructed environment. Call reset() before use.

Return type:

MultiNavigator

Initialize the environment with random positions and objectives.

Parameters:

env (Environment) – The current environment.
key (ArrayLike) – JAX random number generator key.

Returns:

The initialized environment.

Return type:

Environment

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

Advance one step. Actions are forces. The step also applies drag -friction * vel.

Parameters:

env (Environment) – The current environment.
action (jax.Array) – The per-agent action vectors.

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]#

Return the 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]#

Return whether the episode has ended.

The episode ends when step_count exceeds max_steps.

Parameters:: env (Environment) – The current environment.
Returns:: A bool that is True when 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 acts directly on a sphere on a \(z=0\) floor. Each step applies translational drag -friction * vel and angular damping -friction * ang_vel. The environment samples objectives and assigns them one-to-one with 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, and \(K\) is the translational kinetic energy. ke_tau sets the overall strength of the KE penalty. ke_gate controls how sharply KE sensitivity falls off with distance. A 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`
Angular velocity	`3`
LiDAR proximity (normalized)	`n_lidar_rays`

For realistic training parameters, skip_frames = 50 gives a response rate of 200 Hz, so 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:

The constructed environment. Call reset() before use.

Return type:

MultiRoller

Initialize the environment with random positions and objectives.

Parameters:

env (Environment) – The current environment.
key (ArrayLike) – JAX random number generator key.

Returns:

The initialized environment.

Return type:

Environment

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

Advance one step. Actions are torques. The step also applies translational drag and angular damping.

Parameters:

env (Environment) – The current environment.
action (jax.Array) – The per-agent torque vectors.

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,)).
Angular velocity (shape (3,)).
LiDAR proximity, normalized by lidar_range (shape (n_lidar_rays,)).

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

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

Return the 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]#

Return whether the episode has ended.

The episode ends when step_count exceeds max_steps.

Parameters:: env (Environment) – The current environment.
Returns:: A bool that is True when 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.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 acts directly on a sphere inside a reflective box. Each step adds viscous drag -friction * vel. 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 and \(K\) is the translational kinetic energy. ke_tau is the KE scale that sets the overall strength of the penalty. ke_gate controls how sharply KE sensitivity falls off with distance. A 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 penalized 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`

For realistic training parameters, skip_frames = 50 gives a response rate of 200 Hz, so 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.
near_goal_bonus (float) – Reward bonus applied when the agent is within one radius of the objective.
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:

The constructed environment. Call reset() before use.

Return type:

SingleNavigator

Place the agent and the objective at random positions in the box.

Parameters:

env ('SingleNavigator') – The current environment.
key (jax.random.PRNGKey) – JAX random number generator key.

Returns:

The initialized environment.

Return type:

Environment

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

Advance one step. Actions are forces. The step also applies drag -friction * vel.

Parameters:

env (Environment) – The current environment.
action (jax.Array) – The per-agent action vectors.

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]#

Return the 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]#

Return whether the episode has ended.

The episode ends when step_count exceeds max_steps.

Parameters:: env (Environment) – The current environment.
Returns:: A bool that is True when 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 through torque-controlled rolling.

The agent is a sphere resting on a \(z = 0\) floor under gravity. Actions are 3-D torque vectors. Translational motion comes from frictional contact with the floor (see frictional_wall_force()). Each step applies a viscous drag -friction * vel and an angular damping -friction * ang_vel.

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 and \(K\) is the total (translational + rotational) kinetic energy. ke_tau is the KE scale that sets the overall strength of the penalty. ke_gate controls how sharply KE sensitivity falls off with distance. A 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

For realistic training parameters, skip_frames = 50 gives a response rate of 200 Hz, so 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) – Damping coefficient applied as -friction * vel and -friction * ang_vel.
near_goal_bonus (float) – Reward bonus applied when the agent is within one radius of the objective.
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:

The constructed environment. Call reset() before use.

Return type:

SingleRoller

Place the agent and the objective at random positions on the floor.

Parameters:

env (Environment) – The current environment.
key (ArrayLike) – JAX random number generator key.

Returns:

The initialized environment.

Return type:

Environment

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

Apply a torque action and advance the 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]#

Return the 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 that acts on a sphere in a reflective box. Each step adds viscous drag -friction * vel. The environment samples objectives on a jittered grid inside the box. Agents spawn in the padding ring around the box. Each step refreshes three LiDAR sensors: walls, objectives, and peers (other agents). 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 (normalized)	`n_lidar_rays`
Wall LiDAR (normalized)	`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 apparent LiDAR proximity when a peer sits on it. The penalty decays exponentially with the peer-to-objective distance and is zero beyond \(L/4\).

Returns:

The constructed environment. Call reset() before use.

Return type:

SwarmNavigator

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

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

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

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

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

For each objective LiDAR bin, the reward finds the nearest agent over all agent LiDAR bins (distance recovered with the law of cosines). When that agent lies within lr/4 of the objective, the reward subtracts from the objective’s apparent proximity. The penalty decays exponentially and is 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 because 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]#: The episode ends when step_count exceeds max_steps.

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 that acts on a sphere on a \(z=0\) floor. Each step adds translational drag -friction * vel and angular damping -friction * ang_vel. The environment samples objectives on a jittered grid inside the box at floor level. Agents spawn in the padding ring around the box. Each step refreshes three LiDAR sensors: walls, objectives, and peers (other agents). 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 (normalized)	`n_lidar_rays`
Wall LiDAR (normalized)	`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 apparent LiDAR proximity when a peer sits on it. The penalty decays exponentially with the peer-to-objective distance and is zero beyond \(L/4\).

Returns:

The constructed environment. Call reset() before use.

Return type:

SwarmRoller

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

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

static observation(env: SwarmRoller) → Array[source]#: Velocity + angular velocity + objective LiDAR + wall LiDAR (all normalized), 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]#: The episode ends when step_count exceeds max_steps.

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.

Same structure as SwarmRoller: rolling-sphere agents with translational and angular drag, three LiDAR sensors (walls, objectives, peers), and a bin-wise contention-shaped reward. Two differences: the objectives form a square pyramid sensed with 3-D LiDAR, and the agents attract each other pairwise through a magnetic force.

Feature	Size
Velocity	`dim`
Angular velocity	`dim`
Objective LiDAR (normalized)	`n_az * n_el`
Wall LiDAR (normalized)	`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]#: Initialize 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 and sensing.

static observation(env: SwarmRoller3D) → Array[source]#: Velocity + angular velocity + objective LiDAR + wall LiDAR (normalized), 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) because the bins are the flattened 3-D (azimuth, elevation) grid.

static done(env: SwarmRoller3D) → Array[source]#: The episode ends when step_count exceeds max_steps.

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.

The dynamics, pairwise attraction, nearest-neighbor observation, and per-gear reward match TwoGears. Only the objective differs: the num_gears targets form a triangular stack. The rows shrink by one from bottom to top and the gears touch. num_gears=3 gives the classic triangle [2,1], 5 -> [3,2], and 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). It must also be >= 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:

The 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). The “other gear” slot holds the nearest neighbor.

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. A pairwise magnetic force attracts the gears to each other, and each gear observes its nearest neighbor.

Note

The maximum torque must be at least 4.0 * mgr so the gear can climb correctly, and the attraction must be at least 1 * mg. For realistic training parameters, skip_frames = 50 gives a response rate of 200 Hz, so 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:

The 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 current environment.
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.

The 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 holds its nearest neighbor:

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 per-gear reward.

The reward is 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 toward assigned targets.
`multi_roller`	Environment where multiple rolling agents navigate toward assigned targets.
`single_navigator`	Environment where a single agent navigates toward a target.
`single_roller`	Environment where a single agent rolls toward 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`	Environment where N dynamic gears assemble a triangular stack.
`two_gears`	Two-dimensional environment where N dynamic gears assemble a tower.