jaxdem.rl.environments.three_gears#

Three-gear environment: three dynamic gears must assemble a triangle.

Classes

ThreeGears(state, system, env_params, num_gears)

N dynamic gears that must assemble a triangular stack.
class jaxdem.rl.environments.three_gears.ThreeGears(state: State, system: System, env_params: dict[str, Any], num_gears: int)#

Bases: Environment

N dynamic gears that must assemble a triangular stack.

Identical dynamics, pairwise attraction, nearest-neighbour observation, and per-gear reward as TwoGears — only the objective differs: the num_gears targets form a triangular stack (rows that shrink by one from bottom to top, gears touching). num_gears=3 is the classic triangle [2,1]; 5 -> [3,2]; 6 -> [3,2,1]. Gear i is paired with objective i.

Note

As with TwoGears, skip_frames = 50 gives a 200 Hz response rate, so num_steps_epoch = 100 is a 0.5 s horizon. box_size must fit the stack — width 2*m*rr and height (2 + (m-1)*sqrt(3))*rr (m = bottom row size), and >= 2*rr*(num_gears+1) wide for a non-overlapping spawn.

num_gears: int#

Number of gears (agents) forming the triangular stack.

classmethod Create(num_gears: int = 6, box_size: float = 30.0, max_steps: int = 100000, friction: float = 0.2, ke_weight: float = 0.1, attraction_mag: float = 2.0) → ThreeGears[source]#

Create an N-gear triangular-stack environment.

Parameters:

  • num_gears (int) – Number of dynamic gears (agents) forming the stack.

  • box_size (float) – Size of the square bounding box.

  • max_steps (int) – Episode length in physics steps.

  • friction (float) – Viscous drag coefficient applied as -friction * vel.

  • ke_weight (float) – Weight for the differential kinetic energy penalty.

  • attraction_mag (float) – Magnitude of the pairwise attraction force between gears.

Returns:

A freshly constructed environment (call reset() before use).

Return type:

ThreeGears

static reset(env: ThreeGears, key: Array) → Environment[source]#

Reset with the gears on the floor and a random triangular-stack objective.

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

Advance one step: per-gear torque, pairwise attraction, viscous drag.

Attraction on gear \(i\) from gear \(j\) is \(-(C/d_{ij}^3)\,\hat{n}_{ij}\) when \(d_{ij} < 3r\), with \(\hat{n}_{ij}=\mathrm{unit}(\mathbf{r}_i-\mathbf{r}_j)\) and \(C = m_{\text{attr}}(2r)^3\). Net force on \(i\) is \(\sum_{j\ne i}\).

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

Per-gear observation (16 features); “other gear” = nearest neighbour.

Feature

Size

Distance to floor

1

Distance to left/right walls

2

Unit vector to target

2

Clamped displacement to target

2

Unit vector to nearest gear

2

Clamped displacement to nearest gear

2

\(\sin(\Delta\theta)\)

1

\(\cos(\Delta\theta)\)

1

Velocity (x, y)

2

Angular velocity

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

Per-gear shaping reward.

\[R_i = (d_{i,t-1} - d_{i,t}) - w_{\text{ke}} (K_{i,t} - K_{i,t-1})\]
static done(env: ThreeGears) → Array[source]#
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).

property max_num_agents: int[source]#

Maximum number of active agents in the environment.