jaxdem.rl.environments.swarm_roller_3d#

3-D swarm rolling agents covering pyramid objectives, with mutual attraction.

Classes

SwarmRoller3D(state, system, env_params, ...)

Multi-agent cooperative coverage of 3-D pyramid objectives with attraction.

class jaxdem.rl.environments.swarm_roller_3d.SwarmRoller3D(state: State, system: System, env_params: dict[str, Any], n_lidar_rays: int, n_lidar_elevation: int, num_objectives: int)#

Bases: Environment

Multi-agent cooperative coverage of 3-D pyramid objectives with attraction.

Identical in structure to SwarmRoller: rolling-sphere agents with translational and angular drag, three LiDAR sensors (walls, objectives, peers) and a bin-wise contention-shaped reward. Two differences: objectives are arranged as a square pyramid (sensed with 3-D LiDAR), and agents exert pairwise magnetic attraction on each other.

Feature	Size
Velocity	`dim`
Angular velocity	`dim`
Objective LiDAR (normalised)	`n_az * n_el`
Wall LiDAR (normalised)	`n_az * n_el`

n_lidar_rays: int#: Number of azimuthal bins for each 3-D LiDAR sensor.

n_lidar_elevation: int#: Number of elevation bins for each 3-D LiDAR sensor.

num_objectives: int#: Number of objectives (pyramid spheres) sampled per environment.

classmethod Create(N: int = 5, num_objectives: int = 5, box_size: float = 5.0, box_padding: float = 5.0, max_steps: int = 10000, friction: float = 0.2, near_goal_bonus: float = 0.01, lidar_range: float = 16.0, n_lidar_rays: int = 8, n_lidar_elevation: int = 8, contention_strength: float = 15.0, magnet_strength: float = 4.0, magnet_range: float = 3.0) → SwarmRoller3D[source]#

Create a 3-D swarm roller environment with pyramid objectives.

Parameters mirror SwarmRoller.Create(), plus n_lidar_elevation (3-D LiDAR elevation bins) and magnet_strength / magnet_range for the inter-agent attraction.

static reset(env: SwarmRoller3D, key: Array | ndarray | bool | number | bool | int | float | complex) → Environment[source]#: Initialise with agents in the padding ring and a pyramid of objectives in the box.

static step(env: SwarmRoller3D, action: Array) → Environment[source]#: Advance one step: drag, torque, mutual attraction, then physics + sensing.

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

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

Potential-based shaping with a bin-wise contention penalty.

Same as SwarmRoller.reward(), but the law-of-cosines bin geometry uses azimuth alignment (az = bin // n_elevation) since the bins are the flattened 3-D (azimuth, elevation) grid.

static done(env: SwarmRoller3D) → 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.