jaxdem.rl.environments.granulabot#

Granulabots with internal magnetic rotors and TwoGears’ tower task.

Body radius = robot mass = gravity = 1 (24 mm, 98 g). Each 260-sphere, 60-tooth chassis carries active +z and passive -z magnetic rotors, orbiting at 0.75. Their solid-cylinder force profile has radius 0.25, length 15.5/24 and axial spacing 16.5/24; reference Br = 1.45 T. The force uses a beta/six-term blend with default multiplier 0.25, giving about 2mg of attraction between vertically touching robots with their same-polarity magnets facing each other. Mechanical rotor masses and inertias retain the calibrated 98 g robot model. Actions are normalized motor voltages, with 6 V full scale and 20 ms lag. The default motor has twice the reference unloaded speed at the same stall torque; motor_speed_scale=1 restores the paper-calibrated torque-speed law.

Functions

rebuild_status(state, system)

Return geometry, cutoff, and whether cache maintenance may be needed.

Classes

`Granulabot`(state, system, env_params, ...)	TwoGears-style vertical tower with one voltage action per robot.
`MagnetPairs`(pairs, weights, particle_count)	Weighted magnetic interactions between explicit Granulabot particle pairs.

jaxdem.rl.environments.granulabot.rebuild_status(state: State, system: System) → tuple[Array, Array, Array][source]#

Return geometry, cutoff, and whether cache maintenance may be needed.

Geometry changes must reach check_and_rebuild so it can account for image motion before deciding whether the skin budget requires a rebuild.

class jaxdem.rl.environments.granulabot.MagnetPairs(pairs: Array, weights: Array, particle_count: int)#

Bases: BondedForceModel

Weighted magnetic interactions between explicit Granulabot particle pairs.

pairs: Array#

weights: Array#

particle_count: int#

classmethod Create(pairs: Array | ndarray | bool | number | bool | int | float | complex, weights: Array | ndarray | bool | number | bool | int | float | complex, particle_count: int) → MagnetPairs[source]#

property force_and_energy_fns: tuple[Callable[[Array, State, System], tuple[Array, Array]], Callable[[Array, State, System], Array], bool][source]#

Build the bonded force/energy callables for the force manager.

Returns:

(force_fn, energy_fn, is_com_force) where:

force_fn computes bonded force and torque contributions.
energy_fn computes bonded potential-energy contributions.
is_com_force tells where the force acts: True for the center of mass, False for the contact point. It has no effect on spheres.

Return type:

Tuple[ForceFunction, EnergyFunction, bool]

static force(pos: Array, state: State, system: System) → tuple[Array, Array][source]#

static energy(pos: Array, state: State, system: System) → Array[source]#

compute_potential_energy(pos: Array, state: State, system: System) → Array[source]#

static merge(first: BondedForceModel, others: BondedForceModel | Sequence[BondedForceModel]) → MagnetPairs[source]#

Merge two or more bonded-force models into one.

The merge concatenates topology, reference, and coefficient arrays. It shifts vertex indices and body IDs automatically so references stay consistent. When one side has a term that the other does not, the merge pads missing coefficients with 0 and missing reference values with 1.

Parameters:

model1 (BondedForceModel) – Base model.
model2 (BondedForceModel or Sequence[BondedForceModel]) – Model(s) to merge into model1.

Returns:

A new model containing all bodies from both sides.

Return type:

BondedForceModel

Create a new body from raw arrays and merge it into an existing model.

This is a convenience wrapper equivalent to calling the concrete Create constructor followed by merge().

Parameters:

model (BondedForceModel) – Existing model to extend.
**kwargs – Constructor arguments forwarded to the concrete Create method (e.g. vertices, elements, coefficients, …).

Returns:

The extended model.

Return type:

BondedForceModel

class jaxdem.rl.environments.granulabot.Granulabot(state: State, system: System, env_params: dict[str, Array], num_gears: int, neighbor_capacity: int)#

Bases: Environment

TwoGears-style vertical tower with one voltage action per robot.

All robots spawn on the floor; robot i targets tower level i at a random shared x coordinate. Observations include chassis/neighbor geometry and both robots’ rotor and motor states. Actions remain normalized motor voltages, not direct chassis torques. The two internal rotors keep their constrained mechanics.

attraction_mag scales the signed solid-cylinder magnetic interaction. Its default 0.25 gives about 2mg net attraction at a chassis separation of 2.06 radii with active magnets facing and passive magnets facing away. This is an aligned static force calibration, not an all-orientation hold. motor_speed_scale scales unloaded speed at fixed stall torque (default 2). This is a faster motor model, not a change in timestep or playback rate. neighbor_capacity controls the average contact-pair storage per particle; increase it if the collider reports overflow for unusually dense scenes.

num_gears: int#

neighbor_capacity: int#

classmethod Create(num_gears: int = 2, box_size: float = 20.0, max_steps: int = 100000, ke_weight: float = 0.1, attraction_mag: float = 0.25, dt: float = 0.0005, neighbor_capacity: int = 32, motor_speed_scale: float = 2.0) → Granulabot[source]#

static reset(env: Granulabot, key: Array | ndarray | bool | number | bool | int | float | complex) → Granulabot[source]#

static step(env: Granulabot, action: Array) → Granulabot[source]#: Apply voltage and advance physics; preserve the action-start baseline.

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

Save per-robot distance and kinetic energy before the next action.

Parameters:

env (Granulabot) – The environment at the start of the action interval.
action (jax.Array) – The per-agent action about to be applied. Wrappers pass the same transformed action to this hook and to each physics step. This implementation does not require the action value.

Returns:

The environment with prev_dist and prev_ke captured from the live state. Observation and reward accessors compute current values.

Return type:

Environment

static observation(env: Granulabot) → Array[source]#: Compute per-robot observations from the live physical state.

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

static truncated(env: Granulabot) → 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.