jaxdem.rl.environments.swarm_navigator#

Environment where multiple agents cooperatively cover a set of objectives.

Classes

SwarmNavigator(state, system, env_params, ...)

Multi-agent cooperative objective coverage with local sensing.

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