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, n_lidar_elevation: int, num_objectives: int, vacancy_aware_lidar: bool = False, approach_kernel: str = 'legacy')#

Bases: Environment

Multi-agent cooperative objective coverage with local sensing.

Each agent controls a force vector that acts on a sphere in a reflective 2-D or 3-D box. Each step adds viscous drag -friction * vel. Objectives lie on a jittered grid inside the box; agents spawn in its padding ring or shell. Observations read live objective and agent/wall LiDAR. By default, observations hide nearby objectives estimated to be covered by another agent, preserving the observing agent’s own goal. Reward still uses the unfiltered scans and shares value among nearby claimants, combining flat settling with a wider approach potential. Local objective and agent/wall LiDAR estimate peer-to-objective distances across every pair of angular bins. prev_potential holds the action-start potential, captured once by checkpoint and preserved through all physics steps. Reset initializes this baseline. Only the historical potential is stored; current sensor values are computed when observations or rewards are read.

The shared potential combines the best settling value with the summed approach value of visible objectives:

\[\Phi_i(s)=\max_{k\in V_i}\bigl(w_{ik}B(d_{ik})\bigr) +\eta\sum_{k\in V_i}w_{ik}P(d_{ik}/a_i),\qquad r_{i,t}=\Phi_i(s_{t+1})-\Phi_i(s_{\mathrm{checkpoint}}).\]

Here \(a_i\) is the particle radius, \(V_i\) contains raw objective detections, \(B(d)=\exp[-(d/(1.5a_i))^4]\), and \(w_{ik}=1/(1+C_{ik})\) shares value with estimated peer claims \(C_{ik}\). See reward() for the approach potential and references. Objectives have no assigned owners; “own goal” in the observation filter means a detected goal within the observing agent’s radius. Agent and objective counts may differ.

Notes

The observation vector per agent is:

Feature	Size
Velocity	`dim`
Normalized objective LiDAR	`lidar_bins`
Normalized agent/wall LiDAR	`lidar_bins`

Velocity is not normalized. LiDAR stores proximity, with zero for empty or masked bins and one for coincident centers after normalization. With two spatial dimensions and the default 12 bins, each observation has 26 entries. In 3-D, lidar_bins = n_lidar_rays * n_lidar_elevation; default settings give 291 entries. Methods describe one environment; use jax.vmap for independent environment batches.

With dt = 0.002 and skip_frames = 49, one policy action spans 50 physics steps (0.1 seconds). jaxdem.utils.advance_action checkpoints once before these steps. Read reward at their endpoint, before another checkpoint or reset replaces the baseline. Forces are applied directly with viscous drag; this environment has no action smoothing state.

Examples

>>> import jax
>>> import jax.numpy as jnp
>>> from jaxdem.rl import Environment
>>> from jaxdem.utils import advance_action
>>> env = Environment.create("swarmNavigator")
>>> env = env.reset(env, jax.random.key(0))
>>> action = jnp.zeros((env.max_num_agents, env.action_space_size))
>>> env, terminated, truncated = advance_action(env, action, skip_frames=49)
>>> reward = env.reward(env)
>>> observation = env.observation(env)

n_lidar_rays: int#: Number of angular bins for each LiDAR sensor.

n_lidar_elevation: int#: Number of elevation bins in 3-D; unused in 2-D.

num_objectives: int#: Number of objectives sampled per environment.

vacancy_aware_lidar: bool#: Opt in to removing occupied goals before objective LiDAR bin reduction.

approach_kernel: str#: Select the original attraction or an opt-in raised-cosine approach.

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, lidar_range: float = 16.0, n_lidar_rays: int = 12, attraction_coeff: float = 0.1, attraction_constant: float = 0.01, attraction_quadratic: float = 0.02, attraction_decay: float = 0.3333333333333333, sharing_range: float = 1.5, hide_occupied_objectives: bool = True, dim: int = 2, n_lidar_elevation: int = 12, *, vacancy_aware_lidar: bool = False, approach_kernel: str = 'legacy', occupancy_radius_scale: float = 1.2) → 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 or cube that holds the objectives.
box_padding (float) – Total extra domain side length in particle radii. The spawn ring or shell extends by half this amount on each side of the objective box.
max_steps (int) – Episode length in physics steps.
friction (float) – Viscous drag coefficient applied as -friction * vel.
lidar_range (float) – Maximum detection range \(L\) for the LiDAR sensors.
n_lidar_rays (int) – Number of angular LiDAR bins spanning \([-\pi, \pi)\).
attraction_coeff (float) – Nonnegative weight of the summed approach potential, default 0.1. Peer claims reduce both approach and settling values. Zero disables only the approach term.
attraction_constant (float) – Nonnegative coefficients \(c,b\) of the attraction strength \((c+b x^2)e^{-\lambda x}\), with \(x=d/\mathrm{rad}_i\), used by the legacy approach kernel.
attraction_quadratic (float) – Nonnegative coefficients \(c,b\) of the attraction strength \((c+b x^2)e^{-\lambda x}\), with \(x=d/\mathrm{rad}_i\), used by the legacy approach kernel.
attraction_decay (float) – Positive exponential decay \(\lambda\) in radius units. Used by the legacy approach kernel.
sharing_range (float) – Positive quartic claim scale in observing-agent radii. The default matches the settling scale, 1.5.
hide_occupied_objectives (bool) – Enabled by default. Hide readings strictly below half the LiDAR range when a detected peer appears within occupancy_radius_scale observing-agent radii of the objective. Preserve goals within the observer’s radius. A hidden reading becomes zero; farther goals are not revealed. Applies only to observations, never reward or checkpoints. Set to False to recover the unfiltered observation.
vacancy_aware_lidar (bool) – Disabled by default. When enabled, remove goals occupied by visible peers before nearest-per-bin reduction, revealing farther vacant goals in the same bin. Preserve goals covered by the observer. Requires all goal and peer detections within the observer’s LiDAR range before binning; rewards still use the original raw scans. hide_occupied_objectives=False disables either occupancy mask.
dim (int) – Spatial dimension, either 2 (default) or 3. Three-dimensional agents and objectives occupy a volume, with agents in its padding shell.
n_lidar_elevation (int) – Positive number of elevation bins in 3-D; unused in 2-D.
approach_kernel ({"legacy", "cosine"}) – The default retains the original approach potential. "cosine" replaces only that term with (1 + cos(pi*d/lidar_range)) / 2 for visible goals, preserving settling, sharing, observations, and potential-difference timing. This option uses the existing attraction weight and sharing range; the constant, quadratic, and decay attraction controls have no effect on its score.
occupancy_radius_scale (float) – Positive peer-to-objective masking tolerance in observing-agent radii, default 1.2. Set to 1.0 for the original cutoff. Applies to the default post-bin mask only, not vacancy_aware_lidar.

Returns:

The constructed environment. Call reset() before use.

Return type:

SwarmNavigator

Raises:

ValueError – If an attraction coefficient is nonfinite or negative, or attraction_decay, sharing_range, or occupancy_radius_scale is nonfinite or nonpositive, or approach_kernel is not "legacy" or "cosine".

Notes

Construction allocates placeholder agent and objective arrays. Reset samples the geometry and initializes the physical system and reward baseline. N, num_objectives, and n_lidar_rays determine array shapes. Kernel and sensing selections are static fields. Physical, reward, and observation coefficients are stored in env_params as JAX arrays.

Sample agents in the padding ring and objectives in the inner box.

Parameters:

env (SwarmNavigator) – The environment to initialize, retaining its configured sizes and reward/observation parameters.
key (ArrayLike) – JAX random number generator key, split for agents and objectives.

Returns:

The initialized environment with new physical state, zero initial velocities, restarted step count, and an action-start potential equal to the initial potential. Reward is initially zero.

Return type:

Environment

Notes

Agents have radius one. Sampling uses jittered cells with requested center separation 2.05; undersized cells cannot guarantee this spacing. The reflecting domain has side length box_size + box_padding and lower corner -box_padding / 2 in each coordinate. Reset constructs elastic material with unit-mass density (1 / pi in 2-D, 3 / (4*pi) in 3-D), Young’s modulus 2e5, and Poisson ratio 0.3. The physics time step is 0.002 and rotation is disabled. Objective coordinates have shape (num_objectives, dim) and are stored in env_params["objective"].

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

Save the LiDAR-derived potential at the start of a policy action.

Physics steps preserve this potential through all skipped frames. The baseline is independent of action. Read reward before the next checkpoint or reset replaces it. Unchanged states give zero reward.

Parameters:

env (SwarmNavigator) – Environment immediately before the next policy action.
action (jax.Array) – Per-agent force requests, normally shape (N, 2). Accepted for the environment checkpoint interface; its value is unused.

Returns:

Environment with env_params["prev_potential"] replaced by the current per-agent potential of shape (N,). Physical state and observations are not advanced.

Return type:

Environment

static step(env: SwarmNavigator, action: Array) → Environment[source]#

Advance one physics step with the requested force and viscous drag.

The added force on agent \(i\) is

\[\mathbf{f}_i=\mathbf{a}_i-\mu\mathbf{v}_i,\]

where \(\mu\) is friction. The physics system adds contact forces and advances the state. Requests are not smoothed or clipped by this method.

Parameters:

env (SwarmNavigator) – Current environment with an initialized physical system.
action (jax.Array) – Force requests, normally shape (N, 2), reshaped internally to (N, *action_space_shape).

Returns:

Environment after one physics step. The action-start potential is preserved and no sensor readings are cached.

Return type:

Environment

Notes

Use jaxdem.utils.advance_action to checkpoint once before repeating steps for a policy action. Calling this method directly neither creates a checkpoint nor handles episode reset.

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

Return velocity, filtered objective LiDAR, and agent/wall LiDAR.

Parameters:: env (SwarmNavigator) – Environment whose current state is observed. Checkpointed potentials are not used to construct the observation.
Returns:: Shape (N, observation_space_size). Each row concatenates the velocity components, normalized objective proximity bins, then normalized agent/wall proximity bins. Velocity is in simulation units; LiDAR channels lie in [0, 1].
Return type:: jax.Array

Notes

A valid raw LiDAR return is normalized as \(p_k=\max(0,1-d_k/L)\). Empty and masked bins are zero. hide_occupied_objectives=False returns unfiltered goal readings. Only LiDAR channels are normalized. With filtering enabled, hide objective bin \(k\) when

\[a_i < d_{ik} < L/2,\qquad \min_{q\in Q_i}\widehat\delta_{ikq}\le \sigma a_i,\]

where \(a_i\) is the observing radius, \(L\) the LiDAR range, and \(Q_i\) the valid detected peers, excluding walls and self. occupancy_radius_scale sets \(\sigma\) (default 1.2). Bin-center directions give the approximate separation

\[\widehat\delta_{ikq}^2=\max\!\left( d_{ik}^2+d_{iq}^2-2d_{ik}d_{iq}\cos(\theta_k-\theta_q),0\right).\]

Every objective bin is compared with every peer bin. Own goals within one radius and readings at or beyond half range remain visible. Hidden readings become zero without revealing occluded goals. Reward retains the raw scans. Angular quantization and hidden peers can cause occupancy errors; this is not an exact global coverage test. In 3-D, the cosine is the dot product of spherical bin-center unit directions, including elevation.

The half-range mask and own-goal exception were developed for this environment’s experiments. The congestion-game and reward-shaping references in reward() motivate the reward, not this observation filter, and do not guarantee that discarding observations is harmless.

With vacancy_aware_lidar=True, occupied objectives are removed before nearest-per-bin reduction over all local detections instead. Visible peers within range can mask goals at any distance below range; occupancy uses each peer’s own radius and preserves the observer’s covered goals. This can reveal farther vacant goals behind occupied ones, while retaining the same observation layout. Setting hide_occupied_objectives=False disables occupancy filtering.

static reward(env: SwarmNavigator) → Array[source]#

Return the action-start difference in the shared coverage potential.

Parameters:: env (SwarmNavigator) – Environment at the endpoint of the action interval, with prev_potential captured by checkpoint() or reset().
Returns:: Per-agent rewards of shape (N,). Positive values indicate increased potential; negative values indicate decreased potential. Reading reward does not update the baseline or physical state.
Return type:: jax.Array

Notes

Reward uses raw objective detections regardless of the observation mask. It is evaluated over the complete action interval, not summed once per skipped physics frame. There is no separate occupancy bonus.

Normalized goal distance is \(x=d/a_i\), where \(a_i\) is the observing agent’s radius. Attraction strength and its integrated potential are

\[A(x)=(c+bx^2)e^{-\lambda x},\qquad P(x)=e^{-\lambda x}\left[\frac{c}{\lambda} +b\left(\frac{x^2}{\lambda}+\frac{2x}{\lambda^2} +\frac{2}{\lambda^3}\right)\right],\qquad -P'(x)=A(x).\]

attraction_constant, attraction_quadratic, and attraction_decay set \(c,b,\lambda\). Defaults 0.01, 0.02, and 1/3 place the strength peak near 5.92 radii; strength at two radii is about 47 percent of the peak. The potential itself is monotone.

With approach_kernel="cosine", replace this approach potential by \(H(d)=(1+\cos(\pi d/L))/2\) for visible goals, where \(L\) is the sensor range. Its value and slope vanish at \(L\), and the three attraction-shape controls have no effect on it. Settling, soft claim sharing, attraction weight, and reward timing are unchanged. This scalar potential adapts the exploration kernel of Sun et al. (2023), rather than implementing their mean-shift velocity controller.

The potential estimates peer claims on each visible objective from all pairs of objective and peer LiDAR bins:

\[\begin{split}\begin{aligned} C_{ik} &= \sum_{q\in Q_i} \exp\!\left[-\left(\frac{\widehat\delta_{ikq}}{\rho a_i}\right)^4\right], \\ w_{ik} &= \frac{1}{1+C_{ik}}, \\ B_{ik} &= \exp\!\left[-\left(\frac{d_{ik}}{1.5a_i}\right)^4\right], \\ \Phi_i &= \max_{k\in V_i}(w_{ik}B_{ik}) +\eta\sum_{k\in V_i}w_{ik}P(d_{ik}/a_i), \\ r_{i,t} &= \Phi_i(s_{t+1})-\Phi_i(s_{\mathrm{checkpoint}}). \end{aligned}\end{split}\]

sharing_range sets \(\rho\) (default 1.5) and attraction_coeff sets \(\eta\) (default 0.1). The denominator includes one prospective claim for self, so arrival never suppresses the agent’s own value. An objective with no detected claimants has full value; \(n-1\) unit peer claims divide both terms by \(n\). Settling uses the best shared score, rather than an unweighted nearest goal or a sum of nearby settling bonuses. All weights are positive; there is no contact penalty or negative attraction.

This approximates equal sharing in a congestion game (Rosenthal, 1973; see References). Soft counts, angular quantization, occlusion, approach shaping, and discounted physical motion do not inherit the discrete game’s full-coverage equilibrium guarantee. Empty scans have zero potential. No claim assignments are stored.

Potential-based reward shaping is studied by Ng, Harada, and Russell (1999), with a stochastic-game extension by Lu, Schwartz, and Givigi (2011); see References. Their invariance results concern adding \(\gamma\Phi_i(s')-\Phi_i(s)\) to an existing reward, with suitable terminal conditions. Here the undiscounted difference is itself the reward; for \(\gamma<1\) it is not the discount-correct shaping term. These results do not guarantee coverage or convergence, nor do they apply to the observation filter.

Separation for objective bin \(k\) and valid peer bin \(q\) uses bin-center angles:

\[\widehat{\delta}_{ikq}^2=\max\!\left( d_{ik}^2+d_{iq}^2-2d_{ik}d_{iq}\cos(\theta_k-\theta_q),0\right).\]

All bin pairs are compared, including across the angular wraparound. In 3-D, spherical bin-center direction dot products replace the planar cosine and account for both azimuth and elevation. Walls, self IDs, and empty peer bins contribute zero. No peer positions or other agents’ objective scans are used. A wall can hide a peer in its bin, and bin-center directions approximate the true bearing. \(V_i\) contains nonempty raw objective bins; empty scans have zero potential, including goals exactly at sensor range. Setting attraction_coeff=0 retains only the best shared settling value.

Reset and unchanged states give zero reward to floating-point precision. Undiscounted action rewards telescope across contiguous intervals with matching endpoint checkpoints and fixed parameters; episode resets begin a new baseline.

References

Sun, G. et al. (2023). “Mean-shift exploration in shape assembly of robot swarms.” Nature Communications, 14, 3476. https://doi.org/10.1038/s41467-023-39251-5

Rosenthal, R. W. (1973). “A class of games possessing pure-strategy Nash equilibria.” International Journal of Game Theory, 2, 65-67. https://doi.org/10.1007/BF01737559

Ng, A. Y., Harada, D., and Russell, S. (1999). “Policy invariance under reward transformations: Theory and application to reward shaping.” Proceedings of ICML, 278-287. https://ai.stanford.edu/~ang/papers/shaping-icml99.pdf

Lu, X., Schwartz, H. M., and Givigi, S. N. (2011). “Policy invariance under reward transformations for general-sum stochastic games.” Journal of Artificial Intelligence Research, 41, 397-406. https://doi.org/10.1613/jair.3384

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

Return whether the environment has reached its physics-step limit.

Parameters:: env (SwarmNavigator) – Current environment supplying system.step_count and env_params["max_steps"].
Returns:: Scalar boolean, true when step_count >= max_steps. The limit counts physics steps, including skipped frames, not policy actions. This method does not test coverage or reset the environment.
Return type:: jax.Array

property action_space_size: int[source]#

Flattened action size per agent.

Returns:: Number of force components, equal to state.dim (two or three). Actions passed to step() normally have shape (max_num_agents, action_space_size).
Return type:: int

property action_space_shape: tuple[int][source]#

Unflattened shape of one agent’s force request.

Returns:: (state.dim,), either (2,) or (3,). Used when reshaping the batch of per-agent actions inside step().
Return type:: tuple[int]

property observation_space_size: int[source]#

Flattened observation size per agent.

Returns:: state.dim + 2 * lidar_bins (26 with 2-D default settings). observation() returns an array of shape (max_num_agents, observation_space_size). Filtering changes objective readings but not the observation size or channel order.
Return type:: int