jaxdem.rl.environments.two_gears#
Two-dimensional environment where N dynamic gears assemble a tower.
Functions
|
Classes
|
Torque-driven gears assembling an assigned vertical tower. |
- jaxdem.rl.environments.two_gears.frictional_floor_force(pos: Array, state: State, system: System) Tuple[Array, Array][source]#
- class jaxdem.rl.environments.two_gears.TwoGears(state: State, system: System, env_params: dict[str, Any], num_gears: int, n_lidar_rays: int)#
Bases:
EnvironmentTorque-driven gears assembling an assigned vertical tower.
Gear i targets level i, bottom to top, at a random shared x. All gears start on the floor and attract peers within three body radii. Actions are externally scaled torques; viscous drag acts on translation and rotation. Gear contacts use Cundall–Strack friction.
Observations follow MultiNavigator: objective direction, displacement clipped to [-3, 3], COM velocity, angular velocity, and normalized agent/wall LiDAR. Sensors use one center per gear and the actual floor at y=1, excluding constituent spheres of the same gear.
Reward is the action-interval change in exp(-(distance / (2 * outer_radius))**4 - beta * K / K_ref). K includes the whole gear’s translational and rotational energy. The optional energy factor rewards settling near the target; its default zero coefficient matches MultiNavigator’s distance-only reward. Reset/checkpoint save the starting potential, and physics preserves it. Episodes end at the time limit so agents can keep maintaining a stack.
- num_gears: int#
Number of gears (agents) that must form the tower.
- n_lidar_rays: int#
Number of angular LiDAR bins spanning [-pi, pi).
- classmethod Create(num_gears: int = 3, box_size: float = 20.0, max_steps: int = 100000, friction: float = 0.2, attraction_mag: float = 4.0, lidar_range: float = 10.0, n_lidar_rays: int = 16, kinetic_energy_coeff: float = 0.0, kinetic_energy_scale: float = 1.0) TwoGears[source]#
Create a tower environment; call reset before use.
The square box must fit the non-overlapping floor spawn and the entire tower: box_size >= 2 * outer_radius * (num_gears + 1). friction damps linear and angular motion. attraction_mag is the pair force at a center distance of two body radii. Actions remain raw torques, with their magnitude controlled externally.
lidar_range is a center-to-center or center-to-wall distance. n_lidar_rays sets the number of angular bins. kinetic_energy_coeff is nonnegative, and kinetic_energy_scale is a positive reference energy for the complete rigid gear.
- static reset(env: TwoGears, key: Array) Environment[source]#
Sample the tower location and a non-overlapping floor spawn.
- static step(env: TwoGears, action: Array) Environment[source]#
Apply torque, drag and pair attraction, preserving reward history.
- static checkpoint(env: TwoGears, action: Array) Environment[source]#
Capture the live starting potential once before an action interval.
- static observation(env: TwoGears) Array[source]#
Return goal direction/delta, velocity, spin, and normalized LiDAR.
Layout: direction (2), clipped delta (2), COM velocity (2), angular velocity (1), then n_lidar_rays agent/wall proximities. LiDAR uses gear centers and the floor at y=1. As in MultiNavigator, each angular bin reports its closest detection.
- static reward(env: TwoGears) Array[source]#
Return the potential change across the complete action interval.
- property action_space_size: int[source]#
Flattened action size per agent. Actions passed to
step()have shape(A, action_space_size).
- property action_space_shape: tuple[int][source]#
Original per-agent action shape (useful for reshaping inside the environment).
- property observation_space_size: int[source]#
Flattened observation size per agent.
observation()returns shape(A, observation_space_size).