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 applied to a sphere in a reflective box, with viscous drag -friction * vel added each step. Objectives are sampled on a jittered grid inside the box; 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`
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) → 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 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:

SwarmNavigator

static reset(env: SwarmNavigator, key: Array | ndarray | bool | number | bool | int | float | complex) → Environment[source]#: Initialise the environment with random agents (padding) and objectives (box).

static step(env: SwarmNavigator, action: Array) → Environment[source]#: Advance one step. Actions are forces; drag -friction * vel is added.

static observation(env: SwarmNavigator) → Array[source]#: Velocity + objective LiDAR + wall LiDAR (all normalised), per agent.

static reward(env: SwarmNavigator) → 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: SwarmNavigator) → Array[source]#: Episode terminates when max_steps is reached.

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.