jaxdem.rl.environments.swarm_roller#

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

Classes

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

Cooperative objective coverage with torque-controlled rolling spheres.

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

Cooperative objective coverage with torque-controlled rolling spheres.

Objectives are unassigned points on a jittered grid at floor height; agents spawn in the surrounding padding ring. Reward and objective/peer LiDAR semantics match SwarmNavigator. Rolling dynamics match MultiRoller: gravity, frictional contact with the \(z=0\) floor, normal spring particle contacts, and the naive all-pairs collider. No tangential particle-contact history is stored.

Each physics step smooths the requested torque:

\[\mathbf{u}_{i,t}=\alpha\mathbf{a}_{i,t} +(1-\alpha)\mathbf{u}_{i,t-1},\qquad \boldsymbol{\tau}_{i,t}=\mathbf{u}_{i,t}-\mu\boldsymbol{\omega}_{i,t}, \qquad \mathbf{f}_{i,t}=-\mu\mathbf{v}_{i,t}.\]

action_alpha sets \(\alpha\) (default 0.22); friction sets \(\mu\). The previous applied torque belongs to the immediately preceding physics step, independent of checkpoints. Reset initializes it to zero. A constant request completes 99 percent of its transition in 19 physics steps with the default smoothing coefficient.

The shared coverage potential and per-action reward are

\[\Phi_i(s)=\max_{k\in V_i}(w_{ik}B(d_{ik})) +\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, \(B(d)=\exp[-(d/(1.5a_i))^4]\), and \(w_{ik}=1/(1+C_{ik})\) shares objective value with locally estimated peer claims. reward() documents the approach potential \(P\), claim estimate \(C\), parameters, and references. Empty scans give zero potential. There is no separate occupancy bonus or kinetic-energy term. Agent and objective counts may differ.

By default, observations hide objectives below half the LiDAR range if another detected agent is estimated to cover them. Goals within the observer’s own radius remain visible. Reward always uses the raw scans. prev_potential stores only the action-start baseline; current LiDAR readings are computed from live state without persistent sensor caches.

Notes

The observation vector per agent is:

Feature	Size
Planar velocity	`2`
Angular velocity	`3`
Normalized objective LiDAR	`n_lidar_rays`
Normalized agent/wall LiDAR	`n_lidar_rays`

With the default 12 angular bins, each observation has 29 entries. Velocities are not normalized. As in MultiRoller, the 2-D LiDAR utility bins by azimuth but uses full center distances in three dimensions. All centers initially lie at the same height, so these are planar ranges. If particles lift off the floor, the bin-center occupancy estimate ignores relative elevation and becomes less accurate.

With dt=0.002 and skip_frames=49, a policy action spans 50 physics steps (0.1 seconds). jaxdem.utils.advance_action checkpoints once before these steps; reward is read before the next checkpoint or reset. Reset and unchanged states give zero reward. Methods describe one environment; use jax.vmap for independent batches.

Examples

>>> import jax
>>> import jax.numpy as jnp
>>> from jaxdem.rl import Environment
>>> from jaxdem.utils import advance_action
>>> env = Environment.create("swarmRoller")
>>> 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.

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, action_alpha: float = 0.22, 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) → SwarmRoller[source]#

Create a swarm roller with shared objective value and filtered LiDAR.

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) – Total extra domain side length in particle radii. The spawn ring extends by half this amount on each side of the objective box.
max_steps (int) – Episode length in physics steps.
friction (float) – Translational drag and angular damping coefficient, applied as -friction * vel and -friction * ang_vel.
action_alpha (float) – Fraction of the requested torque applied by the smoothing update each physics step, in [0, 1]. Defaults to 0.22, as in MultiRoller. One disables smoothing.
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\).
attraction_quadratic (float) – Nonnegative coefficients \(c,b\) of the attraction strength \((c+b x^2)e^{-\lambda x}\), with \(x=d/\mathrm{rad}_i\).
attraction_decay (float) – Positive exponential decay \(\lambda\) in radius units.
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 one observing-agent radius of the objective. Preserve goals within the observer’s radius. A hidden reading becomes zero; farther goals behind it are not revealed. Applies only to observations, never reward or checkpoints. Set to False to recover the unfiltered observation.

Returns:

The constructed environment. Call reset() before use.

Return type:

SwarmRoller

Raises:

ValueError – If an attraction coefficient is nonfinite or negative, or attraction_decay or sharing_range is nonfinite or nonpositive; or action_alpha is outside [0, 1].

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. The remaining controls are stored in env_params as JAX arrays.

Sample floor-level agents and objectives and initialize rolling physics.

Parameters:

env (SwarmRoller) – Environment retaining its configured geometry, sensor, actuator, and shared-reward parameters.
key (ArrayLike) – JAX random number generator key, split for agents and objectives.

Returns:

Initialized environment with zero velocities and applied torque, restarted step count, and the initial potential checkpointed. Reward is initially zero.

Return type:

Environment

Notes

Agents have unit radius and mass. Objectives and agents start at z=1. The jittered grid requests center spacing 2.05; undersized cells cannot guarantee that separation. The domain’s horizontal side length is box_size + box_padding; its lower vertical boundary is -1, allowing floor contact before boundary reflection. Reset installs gravity (0, 0, -1), the same frictional floor force as MultiRoller, elastic spring particle contacts, and a naive collider. The physics time step is 0.002. Objective coordinates of shape (num_objectives, 3) are stored in env_params["objective"].

static checkpoint(env: SwarmRoller, 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 (SwarmRoller) – Environment immediately before the next policy action.
action (jax.Array) – Per-agent torque requests, normally shape (N, 3). 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. The immediately preceding applied torque is preserved independently.

Return type:

Environment

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

Smooth the requested torque, then advance physics with damping.

applied_action follows each accepted physics step. prev_potential remains fixed until the next checkpoint or reset.

Parameters:

env (Environment) – The current environment.
action (jax.Array) – The per-agent torque vectors.

Returns:

The updated environment state.

Return type:

Environment

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

Return planar velocity, angular velocity, and filtered local LiDAR.

Parameters:: env (SwarmRoller) – Environment whose current state is observed. Checkpointed potentials are not used to construct the observation.
Returns:: Shape (N, 5 + 2 * n_lidar_rays). Each row concatenates planar velocity (two components), angular velocity (three components), normalized objective bins, then normalized agent/wall bins. Velocities are in simulation units; LiDAR channels lie in [0, 1]. This matches SwarmNavigator sensor semantics and MultiRoller velocity channels.
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 a_i,\]

where \(a_i\) is the observing radius, \(L\) the LiDAR range, and \(Q_i\) the valid detected peers, excluding walls and self. 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.

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.

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

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

Parameters:: env (SwarmRoller) – 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.

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

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

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

Parameters:: env (SwarmRoller) – 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:: Three torque components. step() normally receives shape (max_num_agents, 3).
Return type:: int

property action_space_shape: tuple[int][source]#

Unflattened shape of one agent’s torque request.

Returns:: (3,), used to reshape per-agent torque vectors in step().
Return type:: tuple[int]

property observation_space_size: int[source]#

Flattened observation size per agent.

Returns:: 5 + 2 * n_lidar_rays (29 by default). observation() returns shape (max_num_agents, observation_space_size). Filtering does not change sensor sizes or channel order.
Return type:: int