jaxdem.rl.environments.multi_roller#

Environment where multiple rolling agents navigate toward assigned targets.

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.

Each physics step smooths each requested torque:

\[\mathbf{u}_{i,t} = \alpha\,\mathbf{a}_{i,t} + (1 - \alpha)\,\mathbf{u}_{i,t-1}.\]

Here \(\mathbf{a}_{i,t}\) is the requested torque and \(\mathbf{u}_{i,t}\) is the applied torque before angular damping. action_alpha defaults to 0.22. The previous applied torque is from the immediately preceding physics step, independently of checkpoints. Reset initializes \(\mathbf{u}_{i,0}=\mathbf{0}\).

As in MultiNavigator, reward uses a radius-scaled quartic potential:

\[R_i = 1.5\,\mathrm{rad}_i, \qquad \Phi_i(d,K) = \exp\!\left[-\left(\frac{d}{R_i}\right)^4 - \beta\frac{K}{K_{\mathrm{ref}}}\right], \qquad r_{i,t} = \Phi_i(d_{i,t},K_{i,t}) - \Phi_i(d_{i,\mathrm{checkpoint}},K_{i,\mathrm{checkpoint}}).\]

Distance is measured in 3-D from each particle’s center to its own assigned objective, as in SingleRoller. The potential is flat near the objective and rapidly approaches zero far away. Total kinetic energy is

\[K_i = \tfrac12 m_i\|\mathbf{v}_i\|^2 + \tfrac12\boldsymbol{\omega}_{i,b}^{\top} I_{i,b}\boldsymbol{\omega}_{i,b}.\]

Angular velocity and inertia are expressed in the body’s principal frame. kinetic_energy_coeff sets \(\beta\) (default zero), and kinetic_energy_scale sets \(K_{\mathrm{ref}}\) (default one). At fixed distance, slowing translation or rotation increases the potential near the objective and has little effect far away. At rest, the original flat goal potential is recovered. Unchanged distance and energy give zero reward; there is no per-step motion cost or occupancy bonus. Reset gives zero reward.

Notes

The observation vector per agent is:

Feature	Size
Unit direction to objective	`2`
Clamped displacement	`2`
Velocity	`2`
Angular velocity	`3`
Normalized agent/wall LiDAR	`n_lidar_rays`

With dt = 0.002 and skip_frames = 49, each action spans 50 physics steps (0.1 seconds). One checkpoint captures each agent’s starting potential in prev_potential and distance in prev_dist; physics steps preserve this history. Observations and rewards read the live state. Floor forces use the same implementation as SingleRoller. Particle contacts use normal spring forces and the naive all-pairs collider, as in MultiNavigator. Tangential particle-contact friction and contact history are not used; floor friction still drives rolling.

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, action_alpha: float = 0.22, lidar_range: float = 6.0, n_lidar_rays: int = 16, kinetic_energy_coeff: float = 0.0, kinetic_energy_scale: float = 1.0) → 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.
action_alpha (float) – Fraction of the requested torque applied by the smoothing update each physics step, in [0, 1]. One disables smoothing.
lidar_range (float) – Maximum detection range for the LiDAR sensor.
n_lidar_rays (int) – Number of angular LiDAR bins spanning \([-\pi, \pi)\).
kinetic_energy_coeff (float) – Nonnegative dimensionless strength of the kinetic-energy factor. Zero preserves the distance-only potential.
kinetic_energy_scale (float) – Positive reference kinetic energy in simulation units. The energy factor is exp(-kinetic_energy_coeff * K / kinetic_energy_scale). Energy includes both translation and rotation.

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

Smooth the requested torque, then advance physics with damping.

applied_action follows each accepted physics step. prev_dist and prev_potential remain fixed until the next checkpoint or reset.

Parameters:

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

Returns:

The updated environment state.

Return type:

Environment

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

Save the live distance and full potential before the next action.

prev_potential includes both translational and rotational energy; prev_dist retains the 3-D center distance for inspection. The baseline does not depend on action. Physics steps preserve it so reward covers the complete action, including skipped frames. Read reward before a subsequent checkpoint or reset replaces the baseline. The immediately preceding applied torque is preserved independently.

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

Build per-agent observations and LiDAR from the live physical state.

Contents per agent#

Unit vector from objective to agent in the \(xy\) plane (shape (2,)).
Clamped agent-minus-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 change in each agent’s radius-scaled goal potential.

\[\Phi_i(d,K) = \exp\!\left[-\left( \frac{d}{1.5\,\mathrm{rad}_i}\right)^4 - \beta\frac{K}{K_{\mathrm{ref}}}\right], \qquad r_i = \Phi_i(d_i,K_i) - \Phi_i(d_{i,\mathrm{checkpoint}},K_{i,\mathrm{checkpoint}}).\]

Distances use the live 3-D center positions and assigned objectives. kinetic_energy_coeff sets \(\beta\); kinetic_energy_scale sets \(K_{\mathrm{ref}}\). Energy includes translation and rotation. The checkpoint baseline spans the complete action interval. With the default zero coefficient, reward depends only on distance. Returns an array of shape (N,).

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

Return whether the episode has ended.

The episode ends when step_count reaches 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).