jaxdem.rl.environments#

Reinforcement-learning environment interface.

Classes

Environment(state, system, env_params)

Defines the interface for reinforcement-learning environments.

class jaxdem.rl.environments.Environment(state: State, system: System, env_params: dict[str, Any])#

Bases: Factory, ABC

Defines the interface for reinforcement-learning environments.

Let A be the number of agents (A ≥ 1). Single-agent environments still use A=1.
The environment flattens observations and actions per agent to fixed sizes. Use action_space_shape to reshape inside the environment if needed.

Required shapes

Observation: (A, observation_space_size)
Action (input to step()): (A, action_space_size)
Reward: (A,)
Done: scalar boolean for the whole environment

Todo: - Truncated data field: per-agent termination flag - Render method

Example

To define a custom environment, inherit from Environment. Implement the abstract methods:

>>> @Environment.register("MyCustomEnv")
... @jax.tree_util.register_dataclass
... @dataclass(slots=True)
... class MyCustomEnv(Environment):
...     ...

state: State#: Simulation state.

system: System#: Simulation system configuration.

env_params: dict[str, Any]#: Environment-specific parameters.

classmethod Create(dim: int = 2) → Environment[source]#

Initialize the environment to a valid start state.

Parameters:

env ('MyCustomEnv') – The current environment.
key (jax.random.PRNGKey) – JAX random number generator key.

Returns:

The initialized environment.

Return type:

Environment

Reset the environment when done is True.

When done is True, this method calls the environment’s reset method. When done is False, it returns the current environment unchanged.

Parameters:

env (Environment) – The current environment.
done (jax.Array) – A boolean flag that is True at either termination or truncation.
key (jax.random.PRNGKey) – JAX random number generator key used for the reset.

Returns:

The reset environment when done is True, otherwise the unchanged environment.

Return type:

Environment

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

Advance one physics step using per-agent actions.

Advance only the dynamics; preserve historical reward baselines. Use utils.advance_action to snapshot those baselines once before repeating physics steps. Observations and rewards read the live state.

Parameters:

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

Returns:

The updated environment state.

Return type:

Environment

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

Snapshot historical reward quantities before an action interval.

The default implementation is a no-op. Environments whose rewards compare states across time override this hook to store only the required baseline quantities from the current physical state. Derived current quantities need not be stored: observation and reward accessors can compute them from the live state without modifying the baseline.

Parameters:

env (Environment) – The environment at the start of the next action interval.
action (jax.Array) – The per-agent action about to be applied. Overrides may use it when capturing action-dependent history. Wrappers must pass the same transformed action to this hook and to physics steps.

Returns:

The environment with its historical reward baseline captured.

Return type:

Environment

Notes

utils.advance_action calls this hook exactly once before its physics loop, including when skip_frames=0. Physics steps must preserve the baseline until reward and any truncation bootstrap have been consumed. There is no end-of-action checkpoint. Reset must initialize the baseline from the reset state without requiring a policy action.

The hook does not advance physics or reset episodes. Boundary predicates, observations, and rewards must remain valid for the live physical state. For a difference reward, checkpointing a frozen environment captures its unchanged endpoint and gives zero change on the next evaluation interval. This is a reward-history snapshot, not a disk checkpoint operation.

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

Return per-agent observations derived from the live environment.

This accessor must not require an endpoint checkpoint or change the historical reward baseline.

Parameters:: env (Environment) – The current environment.
Returns:: Per-agent observations, shape (A, observation_space_size).
Return type:: jax.Array

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

Return per-agent rewards for the live state and saved action baseline.

Historical terms compare the live endpoint with quantities saved by checkpoint before the action. Read this accessor before reset or the next checkpoint overwrites that baseline. Accessors must not mutate history; repeated reads of the same state return the same reward.

Parameters:: env (Environment) – The current environment.
Returns:: Per-agent rewards for the current state, shape (A,).
Return type:: jax.Array

static done(env: Environment) → Array[source]#

Return whether the episode terminated or was truncated.

Parameters:: env (Environment) – The current environment.
Returns:: A bool that is True when the episode has ended.
Return type:: jax.Array

Notes

Environment subclasses should override terminated and/or truncated. The trainer reads those causes separately; overriding only done does not signal a boundary to its physics loop. Leave the final state available in step so the trainer can read its reward and bootstrap observation before resetting the episode.

static terminated(env: Environment) → Array[source]#

Return a terminal episode boundary.

Terminal transitions do not bootstrap their value target.

static truncated(env: Environment) → Array[source]#

Return a time-limit or external episode boundary.

Truncated transitions reset the episode and recurrent carry but retain value bootstrapping from the final pre-reset observation, unless the same transition is also terminal. The default is False.

static agent_mask(env: Environment) → Array[source]#

Return active-agent flags with shape (max_num_agents,).

At least one agent must remain active in every environment transition. Observations and rewards, including inactive padding, must be finite.

static info(env: Environment) → dict[str, Any][source]#

Return auxiliary diagnostic information.

The default is an empty dict. Subclasses can override this method to provide environment-specific information.

Parameters:: env (Environment) – The current environment.
Returns:: A dictionary with additional information about the environment.
Return type:: Dict

property num_envs: int[source]#: Number of batched environments.

property max_num_agents: int[source]#: Maximum number of active agents in the environment.

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

Modules

`granulabot`	Granulabots with internal magnetic rotors and TwoGears' tower task.
`multi_navigator`	Environment where multiple agents navigate toward assigned targets.
`multi_roller`	Environment where multiple rolling agents navigate toward assigned targets.
`single_navigator`	Environment where a single agent navigates toward a target.
`single_roller`	Environment where a single agent rolls toward a target on the floor.
`swarm_navigator`	Environment where multiple agents cooperatively cover a set of objectives.
`swarm_roller`	Environment where multiple rolling agents cooperatively cover a set of objectives.
`three_gears`	Environment where N dynamic gears assemble a triangular stack.
`two_gears`	Two-dimensional environment where N dynamic gears assemble a tower.