embodichain.learning

Learning systems.

Currently exposes the rl subpackage for on-policy reinforcement learning.

Submodules

rl	On-policy reinforcement learning pipeline.

Reinforcement Learning

On-policy reinforcement learning pipeline.

Algorithms (PPO/GRPO), rollout buffers, collectors, policy/model builders, and the training entry point; rollout data flows as TensorDict objects.

algo	Algorithm registry and construction helpers (BaseAlgorithm, PPO, GRPO, compute_gae, build_algo).
buffer	On-policy rollout buffer (RolloutBuffer) owning preallocated TensorDict storage, plus minibatch iteration helpers.
collector	Collectors that step vectorized environments and assemble rollout data into a preallocated TensorDict layout.
models	Policy-network registration and model construction (ActorCritic, ActorOnly, MLP, Policy).
train
utils	RL helper utilities: algorithm config, optimizers, and observation helpers.

Classes:

DifferentiableTrainer	Coordinate APG updates and truncated-backpropagation boundaries.
DifferentiableTrainerCfg	Configuration for graph-preserving segmented training.
DifferentiableVecEnv	Batched env that preserves the autograd path through step.
LearningVecEnv	Structural interface shared by lightweight vector environments.

Functions:

build_learning_env(name, *, num_envs, ...)	Build a registered lightweight vector environment.
evaluate_episodes(*, policy, env, ...[, ...])	Evaluate exactly num_episodes completed asynchronous episodes.
get_registered_learning_env_names()	Return registered lightweight environment names.
get_trainer_class(algorithm)	Return the trainer compatible with algorithm.
register_learning_env(name[, factory, override])	Register a lightweight vector-environment factory.

class embodichain.learning.rl.DifferentiableTrainer

Bases: object

Coordinate APG updates and truncated-backpropagation boundaries.

Methods:

__init__(cfg, env, policy, algorithm[, ...])
get_summary()	Return the current in-memory training summary.
load_checkpoint(path)	Restore policy, optimizer, and trainer counters.
save_checkpoint([path])	Save policy, optimizer, and trainer counters.
train(total_timesteps)	Train until at least total_timesteps vector transitions exist.

__init__(cfg, env, policy, algorithm, writer=None, eval_env=None)

get_summary()

Return the current in-memory training summary.

Return type:

dict[str, Any]

load_checkpoint(path)

Restore policy, optimizer, and trainer counters.

Return type:

None

save_checkpoint(path=None)

Save policy, optimizer, and trainer counters.

Return type:

str

train(total_timesteps)

Train until at least total_timesteps vector transitions exist.

Return type:

dict[str, Any]

class embodichain.learning.rl.DifferentiableTrainerCfg

Bases: object

Configuration for graph-preserving segmented training.

Methods:

__init__([segment_length, update_horizon, ...])
copy(**kwargs)	Return a new object replacing specified fields with new values.
replace(**kwargs)	Return a new object replacing specified fields with new values.
to_dict()	Convert an object into dictionary recursively.
validate([prefix])	Check the validity of configclass object.

Attributes:

best_eval_metric
best_eval_mode
checkpoint_dir
deterministic_actions
eval_frequency_steps
eval_seed
experiment_name
num_eval_episodes
save_frequency_updates
segment_length
update_horizon
use_wandb

__init__(segment_length=<factory>, update_horizon=<factory>, deterministic_actions=<factory>, checkpoint_dir=<factory>, experiment_name=<factory>, save_frequency_updates=<factory>, eval_frequency_steps=<factory>, num_eval_episodes=<factory>, eval_seed=<factory>, use_wandb=<factory>, best_eval_metric=<factory>, best_eval_mode=<factory>)

best_eval_metric: str

best_eval_mode: str

checkpoint_dir: str

copy(**kwargs)

Return a new object replacing specified fields with new values.

This is especially useful for frozen classes. Example usage:

@configclass(frozen=True)
class C:
    x: int
    y: int

c = C(1, 2)
c1 = c.replace(x=3)
assert c1.x == 3 and c1.y == 2

Parameters:

• obj (object) – The object to replace.

• **kwargs – The fields to replace and their new values.

Return type:

object

Returns:

The new object.

deterministic_actions: bool

eval_frequency_steps: int

eval_seed: int | None

experiment_name: str

num_eval_episodes: int

replace(**kwargs)

Return a new object replacing specified fields with new values.

This is especially useful for frozen classes. Example usage:

@configclass(frozen=True)
class C:
    x: int
    y: int

c = C(1, 2)
c1 = c.replace(x=3)
assert c1.x == 3 and c1.y == 2

Parameters:

• obj (object) – The object to replace.

• **kwargs – The fields to replace and their new values.

Return type:

object

Returns:

The new object.

save_frequency_updates: int

segment_length: int

to_dict()

Convert an object into dictionary recursively.

Note

Ignores all names starting with “__” (i.e. built-in methods).

Parameters:

obj (object) – An instance of a class to convert.

Raises:

ValueError – When input argument is not an object.

Return type:

dict[str, Any]

Returns:

Converted dictionary mapping.

update_horizon: int | None

use_wandb: bool

validate(prefix='')

Check the validity of configclass object.

This function checks if the object is a valid configclass object. A valid configclass object contains no MISSING entries.

Parameters:

• obj (object) – The object to check.

• prefix (str) – The prefix to add to the missing fields. Defaults to ‘’.

Return type:

list[str]

Returns:

A list of missing fields.

Raises:

TypeError – When the object is not a valid configuration object.

class embodichain.learning.rl.DifferentiableVecEnv

Bases: LearningVecEnv, Protocol

Batched env that preserves the autograd path through step.

detach_state is the truncated-backpropagation boundary: detach differentiable internal state and return the current observation without resetting or resampling the episode. Finished rows must auto-reset inside step, returning the terminal reward/done with the next initial observation.

Methods:

detach_state()

Detach internal state and return its current detached observation.

detach_state()

Detach internal state and return its current detached observation.

Return type:

Tensor | TensorDict

class embodichain.learning.rl.LearningVecEnv

Bases: Protocol

Structural interface shared by lightweight vector environments.

Methods:

__init__(*args, **kwargs)
close()	Release owned resources.
reset(*[, seed, options])	Reset all environments and return the initial observation.
step(action)	Advance all environments by one step.

Attributes:

device
num_envs
single_action_space
single_observation_space

__init__(*args, **kwargs)

close()

Release owned resources.

Return type:

None

device: device

num_envs: int

reset(*, seed=None, options=None)

Reset all environments and return the initial observation.

Return type:

tuple[Tensor | TensorDict, dict[str, Any]]

single_action_space: Space

single_observation_space: Space

step(action)

Advance all environments by one step.

Return type:

tuple[Tensor | TensorDict, Tensor, Tensor, Tensor, dict[str, Any]]

embodichain.learning.rl.build_learning_env(name, *, num_envs, device, **cfg)

Build a registered lightweight vector environment.

Return type:

LearningVecEnv

embodichain.learning.rl.evaluate_episodes(*, policy, env, num_episodes, device, seed=None, on_step=None)

Evaluate exactly num_episodes completed asynchronous episodes.

Return type:

dict[str, float]

embodichain.learning.rl.get_registered_learning_env_names()

Return registered lightweight environment names.

Return type:

list[str]

embodichain.learning.rl.get_trainer_class(algorithm)

Return the trainer compatible with algorithm.

Return type:

type[Trainer] | type[DifferentiableTrainer]

embodichain.learning.rl.register_learning_env(name, factory=None, *, override=False)

Register a lightweight vector-environment factory.

The function supports both @register_learning_env("Name") and direct register_learning_env("Name", Factory) use.

Return type:

Callable[[Callable[..., LearningVecEnv]], Callable[..., LearningVecEnv]] | Callable[..., LearningVecEnv]