jaxdem.rl.environments.multi_roller#

Environment where multiple rolling agents navigate toward assigned targets.

Functions

frictional_wall_force(pos, state, system)

Normal, frictional, and restitution forces for spheres on a \(z = 0\) plane.

Classes

MultiRoller(state, system, env_params, ...)

Multi-agent rolling environment toward assigned targets.

class jaxdem.rl.environments.multi_roller.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)\]

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: 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).