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

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

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