jaxdem.rl.environments.swarm_roller#

Environment where multiple rolling agents cooperatively cover a set of objectives.

Classes

SwarmRoller(state, system, env_params, ...)

Multi-agent cooperative objective coverage with rolling dynamics.

class jaxdem.rl.environments.swarm_roller.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.

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) → SwarmRoller[source]#

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.

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