embodichain.learning.rl

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.

Overview

The embodichain.learning.rl package contains algorithm registries, rollout collection logic, policy/model builders, and training entry points.

Submodules

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.

Top-level APIs

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

Algorithms

Algorithm registry and construction helpers (BaseAlgorithm, PPO, GRPO, compute_gae, build_algo).

Classes:

APG	Optimize policy parameters through differentiable rollout rewards.
APGCfg	Analytic policy-gradient config.
BaseAlgorithm	Base class for RL algorithms.
GRPO	Group Relative Policy Optimization on top of TensorDict rollouts.
GRPOCfg	Configuration for GRPO.
PPO	PPO algorithm consuming TensorDict rollouts.
PPOCfg	Configuration for the PPO algorithm.
RolloutKind	Rollout semantics required by an algorithm.

Functions:

build_algo(name, cfg_kwargs, policy, device, *)
compute_gae(rollout, gamma, gae_lambda)	Compute GAE over a rollout stored as [num_envs, time + 1].
get_registered_algo_names()
segmented_discounted_return(rollout, gamma)	Compute one discounted return per environment within a rollout segment.

class embodichain.learning.rl.algo.APG

Bases: BaseAlgorithm[DifferentiableRollout]

Optimize policy parameters through differentiable rollout rewards.

Methods:

__init__(cfg, policy)
accumulate_segment(rollout)	Accumulate gradients from one TBPTT segment without stepping the optimizer.
begin_update()
cancel_update()
finish_update()	Clip gradients and apply one optimizer step.
update(rollout)	Apply one pathwise-gradient update from a rollout segment.

Attributes:

rollout_kind

__init__(cfg, policy)

accumulate_segment(rollout)

Accumulate gradients from one TBPTT segment without stepping the optimizer.

Return type:

None

begin_update()

Return type:

None

cancel_update()

Return type:

None

device: device

finish_update()

Clip gradients and apply one optimizer step.

Return type:

Dict[str, float]

lr_scheduler: LRScheduler | None

optimizer: Optimizer

rollout_kind = 'differentiable'

update(rollout)

Apply one pathwise-gradient update from a rollout segment.

Return type:

Dict[str, float]

class embodichain.learning.rl.algo.APGCfg

Bases: AlgorithmCfg

Analytic policy-gradient config.

gamma applies within each TBPTT segment and restarts after done.

Methods:

__init__([device, optimizer, lr_scheduler, ...])
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:

batch_size
device
ent_coef
gae_lambda
gamma
lr_scheduler
max_grad_norm
optimizer
skip_nonfinite_updates

__init__(device=<factory>, optimizer=<factory>, lr_scheduler=<factory>, batch_size=<factory>, gamma=<factory>, gae_lambda=<factory>, max_grad_norm=<factory>, ent_coef=<factory>, skip_nonfinite_updates=<factory>)

batch_size: int

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.

device: str

ent_coef: float

gae_lambda: float

gamma: float

lr_scheduler: LRSchedulerCfg

max_grad_norm: float

optimizer: OptimizerCfg

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.

skip_nonfinite_updates: bool

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.

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.algo.BaseAlgorithm

Bases: ABC, Generic[RolloutT]

Base class for RL algorithms.

Methods:

bind_schedule(*, total_updates)	Bind horizon-dependent LR schedules from the training budget.
current_learning_rate()
update(rollout)	Update policy using collected data and return training losses.

Attributes:

device
lr_scheduler
optimizer
rollout_kind

bind_schedule(*, total_updates)

Bind horizon-dependent LR schedules from the training budget.

Return type:

None

current_learning_rate()

Return type:

float

device: device

lr_scheduler: LRScheduler | None

optimizer: Optimizer

rollout_kind = 'standard'

abstract update(rollout)

Update policy using collected data and return training losses.

Return type:

Dict[str, float]

class embodichain.learning.rl.algo.GRPO

Bases: BaseAlgorithm[TensorDict]

Group Relative Policy Optimization on top of TensorDict rollouts.

Methods:

__init__(cfg, policy)
update(rollout)	Update policy using collected data and return training losses.

__init__(cfg, policy)

update(rollout)

Update policy using collected data and return training losses.

Return type:

Dict[str, float]

class embodichain.learning.rl.algo.GRPOCfg

Bases: AlgorithmCfg

Configuration for GRPO.

Methods:

__init__([device, optimizer, lr_scheduler, ...])
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:

batch_size
clip_coef
device
ent_coef
eps
gae_lambda
gamma
group_size
kl_coef
lr_scheduler
max_grad_norm
n_epochs
optimizer
reset_every_rollout
truncate_at_first_done

__init__(device=<factory>, optimizer=<factory>, lr_scheduler=<factory>, batch_size=<factory>, gamma=<factory>, gae_lambda=<factory>, max_grad_norm=<factory>, n_epochs=<factory>, clip_coef=<factory>, ent_coef=<factory>, kl_coef=<factory>, group_size=<factory>, eps=<factory>, reset_every_rollout=<factory>, truncate_at_first_done=<factory>)

batch_size: int

clip_coef: float

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.

device: str

ent_coef: float

eps: float

gae_lambda: float

gamma: float

group_size: int

kl_coef: float

lr_scheduler: LRSchedulerCfg

max_grad_norm: float

n_epochs: int

optimizer: OptimizerCfg

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.

reset_every_rollout: bool

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.

truncate_at_first_done: 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.algo.PPO

Bases: BaseAlgorithm[TensorDict]

PPO algorithm consuming TensorDict rollouts.

Methods:

__init__(cfg, policy)
update(rollout)	Update the policy using a collected rollout.

__init__(cfg, policy)

update(rollout)

Update the policy using a collected rollout.

Return type:

Dict[str, float]

class embodichain.learning.rl.algo.PPOCfg

Bases: AlgorithmCfg

Configuration for the PPO algorithm.

Methods:

__init__([device, optimizer, lr_scheduler, ...])
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:

batch_size
clip_coef
device
ent_coef
gae_lambda
gamma
lr_scheduler
max_grad_norm
n_epochs
optimizer
vf_coef

__init__(device=<factory>, optimizer=<factory>, lr_scheduler=<factory>, batch_size=<factory>, gamma=<factory>, gae_lambda=<factory>, max_grad_norm=<factory>, n_epochs=<factory>, clip_coef=<factory>, ent_coef=<factory>, vf_coef=<factory>)

batch_size: int

clip_coef: float

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.

device: str

ent_coef: float

gae_lambda: float

gamma: float

lr_scheduler: LRSchedulerCfg

max_grad_norm: float

n_epochs: int

optimizer: OptimizerCfg

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.

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.

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.

vf_coef: float

class embodichain.learning.rl.algo.RolloutKind

Bases: str, Enum

Rollout semantics required by an algorithm.

Attributes:

DIFFERENTIABLE
STANDARD

Methods:

__new__(value)

DIFFERENTIABLE = 'differentiable'

STANDARD = 'standard'

__new__(value)

embodichain.learning.rl.algo.build_algo(name, cfg_kwargs, policy, device, *, distributed=False)

embodichain.learning.rl.algo.compute_gae(rollout, gamma, gae_lambda)

Compute GAE over a rollout stored as [num_envs, time + 1].

Parameters:

• rollout (TensorDict) – Rollout TensorDict where value[:, -1] stores the bootstrap value for the final observation and transition-only fields reserve their last slot as padding.

• gamma (float) – Discount factor.

• gae_lambda (float) – GAE lambda coefficient.

Return type:

tuple[Tensor, Tensor]

Returns:

Tuple of (advantages, returns), both shaped [num_envs, time].

embodichain.learning.rl.algo.get_registered_algo_names()

Return type:

list[str]

embodichain.learning.rl.algo.segmented_discounted_return(rollout, gamma)

Compute one discounted return per environment within a rollout segment.

Return type:

Tensor

Environments

Contracts and registration helpers for lightweight learning environments.

Classes:

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.
get_registered_learning_env_names()	Return registered lightweight environment names.
register_learning_env(name[, factory, override])	Register a lightweight vector-environment factory.

class embodichain.learning.rl.env.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.env.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.env.build_learning_env(name, *, num_envs, device, **cfg)

Build a registered lightweight vector environment.

Return type:

LearningVecEnv

embodichain.learning.rl.env.get_registered_learning_env_names()

Return registered lightweight environment names.

Return type:

list[str]

embodichain.learning.rl.env.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]

Evaluation

Shared deterministic episode evaluation for all RL trainers.

Functions:

evaluate_episodes(*, policy, env, ...[, ...])

Evaluate exactly num_episodes completed asynchronous episodes.

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

Evaluate exactly num_episodes completed asynchronous episodes.

Return type:

dict[str, float]

Routing

Trainer routing based on an algorithm’s rollout semantics.

Functions:

get_trainer_class(algorithm)

Return the trainer compatible with algorithm.

embodichain.learning.rl.routing.get_trainer_class(algorithm)

Return the trainer compatible with algorithm.

Return type:

type[Trainer] | type[DifferentiableTrainer]

Differentiable Trainer

Training orchestration for truncated differentiable rollouts.

Classes:

DifferentiableTrainer	Coordinate APG updates and truncated-backpropagation boundaries.
DifferentiableTrainerCfg	Configuration for graph-preserving segmented training.

class embodichain.learning.rl.differentiable_trainer.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.differentiable_trainer.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.

Rollout Buffer

On-policy rollout buffer (RolloutBuffer) owning preallocated TensorDict storage, plus minibatch iteration helpers.

Classes:

RolloutBuffer

Single-rollout buffer backed by a preallocated TensorDict.

Functions:

iterate_minibatches(rollout, batch_size, device)	Yield shuffled minibatches from a flattened rollout.
transition_view(rollout[, flatten])	Build a transition-aligned TensorDict from a rollout.

class embodichain.learning.rl.buffer.RolloutBuffer

Bases: object

Single-rollout buffer backed by a preallocated TensorDict.

The shared rollout uses a uniform [num_envs, time + 1] layout. For transition-only fields such as action, reward, and done, the final time index is reused as padding so the collector, environment, and algorithms can share a single TensorDict batch shape.

Methods:

__init__(num_envs, rollout_len, obs_dim, ...)
add(rollout)	Mark the shared rollout as ready for consumption.
get([flatten])	Return the stored rollout and clear the buffer.
is_full()	Return whether a rollout is waiting to be consumed.
start_rollout()	Return the shared rollout TensorDict for collector write-in.

Attributes:

buffer

__init__(num_envs, rollout_len, obs_dim, action_dim, device)

add(rollout)

Mark the shared rollout as ready for consumption.

Return type:

None

property buffer: TensorDict

get(flatten=True)

Return the stored rollout and clear the buffer.

When flatten is True, the rollout is first converted to a transition view that drops the padded final slot from transition-only fields.

Return type:

TensorDict

is_full()

Return whether a rollout is waiting to be consumed.

Return type:

bool

start_rollout()

Return the shared rollout TensorDict for collector write-in.

Return type:

TensorDict

embodichain.learning.rl.buffer.iterate_minibatches(rollout, batch_size, device)

Yield shuffled minibatches from a flattened rollout.

Return type:

Iterator[TensorDict]

embodichain.learning.rl.buffer.transition_view(rollout, flatten=False)

Build a transition-aligned TensorDict from a rollout.

The shared rollout uses a uniform [num_envs, time + 1] layout. For transition-only fields such as action, reward, and done, the final slot is reserved as padding so that all rollout fields share the same batch shape. This helper drops that padded slot and exposes the valid transition slices as a TensorDict with batch shape [num_envs, time].

Parameters:

• rollout (TensorDict) – Rollout TensorDict with root batch shape [num_envs, time + 1].

• flatten (bool) – If True, return a flattened [num_envs * time] view.

Return type:

TensorDict

Returns:

TensorDict containing transition-aligned fields.

Collectors

Collectors that step vectorized environments and assemble rollout data into a preallocated TensorDict layout.

Classes:

BaseCollector	Base class for rollout collectors.
DifferentiableCollector	Collect graph-preserving rollouts without a preallocated buffer.
DifferentiableRollout	An immutable sequence of graph-preserving transitions.
DifferentiableTransition	One graph-preserving environment transition.
SyncCollector	Synchronously collect rollouts from a vectorized environment.

class embodichain.learning.rl.collector.BaseCollector

Bases: ABC

Base class for rollout collectors.

Methods:

collect(num_steps[, rollout, on_step_callback])

Collect a rollout and return it as a TensorDict.

abstract collect(num_steps, rollout=None, on_step_callback=None)

Collect a rollout and return it as a TensorDict.

Return type:

TensorDict

class embodichain.learning.rl.collector.DifferentiableCollector

Bases: object

Collect graph-preserving rollouts without a preallocated buffer.

Methods:

__init__(env, policy, device)
collect(num_steps, *[, deterministic, ...])	Collect a graph-preserving rollout segment.
detach_state()	Start a new truncated-backpropagation segment.
reset(*[, seed])	Reset the environment and collector state.

__init__(env, policy, device)

collect(num_steps, *, deterministic=False, on_step_callback=None)

Collect a graph-preserving rollout segment.

Parameters:

• num_steps (int) – Number of differentiable environment steps.

• deterministic (bool) – Whether to use deterministic policy actions.

• on_step_callback (Optional[Callable[[DifferentiableTransition], None]]) – Optional callback invoked with each transition.

Return type:

DifferentiableRollout

Returns:

An immutable differentiable rollout.

Raises:

ValueError – If num_steps is not positive.

detach_state()

Start a new truncated-backpropagation segment.

Return type:

Tensor

reset(*, seed=None)

Reset the environment and collector state.

Return type:

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

class embodichain.learning.rl.collector.DifferentiableRollout

Bases: object

An immutable sequence of graph-preserving transitions.

Methods:

__init__(initial_observation, transitions)

Attributes:

final_observation	Return the observation after the final transition.
initial_observation
num_steps	Return the number of collected transitions.
rewards	Stack rewards as [time, num_envs] without detaching them.
transitions

__init__(initial_observation, transitions)

property final_observation: Tensor

Return the observation after the final transition.

initial_observation: Tensor

property num_steps: int

Return the number of collected transitions.

property rewards: Tensor

Stack rewards as [time, num_envs] without detaching them.

transitions: tuple[DifferentiableTransition, ...]

class embodichain.learning.rl.collector.DifferentiableTransition

Bases: object

One graph-preserving environment transition.

Methods:

__init__(observation, policy_output, reward, ...)

Attributes:

action	Return the differentiable policy action.
done	Return the combined termination mask.
info
next_observation
observation
policy_output
reward
terminated
truncated

__init__(observation, policy_output, reward, terminated, truncated, next_observation, info)

property action: Tensor

Return the differentiable policy action.

property done: Tensor

Return the combined termination mask.

info: dict[str, Any]

next_observation: Tensor

observation: Tensor

policy_output: TensorDict

reward: Tensor

terminated: Tensor

truncated: Tensor

class embodichain.learning.rl.collector.SyncCollector

Bases: BaseCollector

Synchronously collect rollouts from a vectorized environment.

Methods:

__init__(env, policy, device[, ...])
collect(num_steps[, rollout, on_step_callback])	Collect a rollout and return it as a TensorDict.

__init__(env, policy, device, reset_every_rollout=False)

collect(num_steps, rollout=None, on_step_callback=None)

Collect a rollout and return it as a TensorDict.

Return type:

TensorDict

Policy Models

Policy-network registration and model construction (ActorCritic, ActorOnly, MLP, Policy).

Classes:

ActorCritic	Actor-Critic with learnable log_std for Gaussian policy.
ActorOnly	Actor-only policy for algorithms that do not use a value function (e.g., GRPO).
MLP	General MLP supporting custom last activation, orthogonal init, and output reshape.
Policy	Abstract base class that all RL policies must implement.

Functions:

build_mlp_from_cfg(module_cfg, in_dim, out_dim)	Construct an MLP module from a minimal json-like config.
build_policy(policy_block, obs_space, ...[, ...])	Build a policy from config using spaces for extensibility.
get_policy_class(name)
get_registered_policy_names()
register_policy(name, policy_cls)

class embodichain.learning.rl.models.ActorCritic

Bases: Policy

Actor-Critic with learnable log_std for Gaussian policy.

This is a placeholder implementation of the Policy interface that: - Encapsulates MLP networks (actor + critic) that need to be trained by RL algorithms - Handles internal computation: MLP output → mean + learnable log_std → Normal distribution - Provides a uniform interface for RL algorithms (PPO, SAC, etc.)

This allows seamless swapping with other policy implementations (e.g., VLAPolicy) without modifying RL algorithm code.

Implements TensorDict-native interfaces while preserving get_action() compatibility for evaluation and legacy call-sites.

Methods:

__init__(obs_dim, action_dim, device, actor, ...)	Initialize internal Module state, shared by both nn.Module and ScriptModule.
evaluate_actions(tensordict)	Evaluate actions and return current policy outputs.
forward(tensordict[, deterministic])	Write sampled actions and value estimates into the TensorDict.
get_differentiable_action(tensordict[, ...])	Sample an action with pathwise gradients.
get_value(tensordict)	Write value estimate for the given observations into the TensorDict.

__init__(obs_dim, action_dim, device, actor, critic)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

evaluate_actions(tensordict)

Evaluate actions and return current policy outputs.

Parameters:

tensordict (TensorDict) – TensorDict containing obs and action.

Return type:

TensorDict

Returns:

A new TensorDict containing sample_log_prob, entropy, and value.

forward(tensordict, deterministic=False)

Write sampled actions and value estimates into the TensorDict.

Return type:

TensorDict

get_differentiable_action(tensordict, deterministic=False)

Sample an action with pathwise gradients.

Return type:

TensorDict

get_value(tensordict)

Write value estimate for the given observations into the TensorDict.

Parameters:

tensordict (TensorDict) – Input TensorDict containing obs.

Return type:

TensorDict

Returns:

TensorDict with value populated.

class embodichain.learning.rl.models.ActorOnly

Bases: Policy

Actor-only policy for algorithms that do not use a value function (e.g., GRPO).

Same interface as ActorCritic: get_action and evaluate_actions return (action, log_prob, value), but value is always zeros since no critic is used.

Methods:

__init__(obs_dim, action_dim, device, actor)	Initialize internal Module state, shared by both nn.Module and ScriptModule.
evaluate_actions(tensordict)	Evaluate actions and return current policy outputs.
forward(tensordict[, deterministic])	Write sampled actions and value estimates into the TensorDict.
get_differentiable_action(tensordict[, ...])	Sample an action with pathwise gradients.
get_value(tensordict)	Write value estimate for the given observations into the TensorDict.

__init__(obs_dim, action_dim, device, actor)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

evaluate_actions(tensordict)

Evaluate actions and return current policy outputs.

Parameters:

tensordict (TensorDict) – TensorDict containing obs and action.

Return type:

TensorDict

Returns:

A new TensorDict containing sample_log_prob, entropy, and value.

forward(tensordict, deterministic=False)

Write sampled actions and value estimates into the TensorDict.

Return type:

TensorDict

get_differentiable_action(tensordict, deterministic=False)

Sample an action with pathwise gradients.

Return type:

TensorDict

get_value(tensordict)

Write value estimate for the given observations into the TensorDict.

Parameters:

tensordict (TensorDict) – Input TensorDict containing obs.

Return type:

TensorDict

Returns:

TensorDict with value populated.

class embodichain.learning.rl.models.MLP

Bases: Sequential

General MLP supporting custom last activation, orthogonal init, and output reshape.

Parameters:

• input_dim (-) – input dimension

• output_dim (-) – output dimension (int or shape tuple/list)

• hidden_dims (-) – hidden layer sizes, e.g. [256, 256]

• activation (-) – hidden layer activation name (relu/elu/tanh/gelu/silu)

• last_activation (-) – last-layer activation name or None for linear

• use_layernorm (-) – whether to add LayerNorm after each hidden linear layer

• dropout_p (-) – dropout probability for hidden layers (0 disables)

Methods:

__init__(input_dim, output_dim, hidden_dims)	Initialize internal Module state, shared by both nn.Module and ScriptModule.
init_orthogonal([scales])	Orthogonal-initialize linear layers and zero the bias.

__init__(input_dim, output_dim, hidden_dims, activation='elu', last_activation=None, use_layernorm=False, dropout_p=0.0)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

init_orthogonal(scales=1.0)

Orthogonal-initialize linear layers and zero the bias.

scales: single gain value or a sequence with length equal to the number of linear layers.

Return type:

None

class embodichain.learning.rl.models.Policy

Bases: Module, ABC

Abstract base class that all RL policies must implement.

A Policy: - Encapsulates neural networks that are trained by RL algorithms - Handles internal computations (e.g., network output → distribution) - Provides a uniform interface for algorithms (PPO, SAC, etc.)

Methods:

__init__()	Initialize internal Module state, shared by both nn.Module and ScriptModule.
evaluate_actions(tensordict)	Evaluate actions and return current policy outputs.
forward(tensordict[, deterministic])	Write sampled actions and value estimates into the TensorDict.
get_action(tensordict[, deterministic])	Sample actions into the provided TensorDict without gradients.
get_differentiable_action(tensordict[, ...])	Sample actions while preserving gradients to policy parameters.
get_value(tensordict)	Write value estimate for the given observations into the TensorDict.

Attributes:

device

Device where the policy parameters are located.

__init__()

Initialize internal Module state, shared by both nn.Module and ScriptModule.

device: torch.device

Device where the policy parameters are located.

abstract evaluate_actions(tensordict)

Evaluate actions and return current policy outputs.

Parameters:

tensordict (TensorDict) – TensorDict containing obs and action.

Return type:

TensorDict

Returns:

A new TensorDict containing sample_log_prob, entropy, and value.

abstract forward(tensordict, deterministic=False)

Write sampled actions and value estimates into the TensorDict.

Return type:

TensorDict

get_action(tensordict, deterministic=False)

Sample actions into the provided TensorDict without gradients.

Parameters:

• tensordict (TensorDict) – Input TensorDict containing obs.

• deterministic (bool) – If True, return the mean action; otherwise sample

Return type:

TensorDict

Returns:

TensorDict with action, sample_log_prob, and value populated.

get_differentiable_action(tensordict, deterministic=False)

Sample actions while preserving gradients to policy parameters.

Stochastic implementations must use a reparameterized sample such as torch.distributions.Distribution.rsample(). The base implementation fails explicitly to prevent a non-reparameterized sample() call from silently producing zero pathwise gradients.

Parameters:

• tensordict (TensorDict) – Input TensorDict containing obs.

• deterministic (bool) – If True, return a differentiable deterministic action.

Return type:

TensorDict

Returns:

TensorDict with differentiable policy outputs populated.

Raises:

NotImplementedError – If the policy has no differentiable sampling implementation.

abstract get_value(tensordict)

Write value estimate for the given observations into the TensorDict.

Parameters:

tensordict (TensorDict) – Input TensorDict containing obs.

Return type:

TensorDict

Returns:

TensorDict with value populated.

training: bool

embodichain.learning.rl.models.build_mlp_from_cfg(module_cfg, in_dim, out_dim)

Construct an MLP module from a minimal json-like config.

Return type:

MLP

Expected schema:

module_cfg = {

“type”: “mlp”, “hidden_sizes”: [256, 256], “activation”: “relu”,

}

embodichain.learning.rl.models.build_policy(policy_block, obs_space, action_space, device, actor=None, critic=None)

Build a policy from config using spaces for extensibility.

Built-in MLP policies still resolve flattened obs_dim / action_dim, while custom policies may accept richer obs_space / action_space inputs.

Return type:

Policy

embodichain.learning.rl.models.get_policy_class(name)

Return type:

Optional[Type[Policy]]

embodichain.learning.rl.models.get_registered_policy_names()

Return type:

list[str]

embodichain.learning.rl.models.register_policy(name, policy_cls)

Return type:

None

Training

Functions:

cli([argv])	Command-line interface for RL training.
parse_args([argv])	Parse command-line arguments.
train_from_config(config_path[, ...])	Run training from a config file path.

embodichain.learning.rl.train.cli(argv=None)

Command-line interface for RL training.

Parses CLI arguments and launches training from a config file.

Task packages are discovered (and init hooks executed) before training so that task environments registered in separate packages (e.g. embodichain_tasks) are available to build_env. This mirrors the run_env CLI.

Return type:

None

embodichain.learning.rl.train.parse_args(argv=None)

Parse command-line arguments.

Parameters:

argv (Sequence[str] | None) – Arguments excluding the command name. Uses sys.argv when omitted.

Return type:

Namespace

Returns:

Parsed training arguments.

embodichain.learning.rl.train.train_from_config(config_path, distributed=None, *, profile=False, profile_output=None)

Run training from a config file path.

Parameters:

• config_path (str) – Path to the training config file (.json, .yaml, or .yml).

• distributed (bool | None) – If True, run multi-GPU distributed training. If None, use trainer.distributed from config.

• profile (bool) – Enable gym EnvProfiler on the training environment.

• profile_output (str | None) – Optional JSON dump path for the profiling report.

Utilities

RL helper utilities: algorithm config, optimizers, and observation helpers.

Classes:

AlgorithmCfg	Shared fields for RL algorithm configs.
LRSchedulerCfg	Optional LR scheduler.
OptimizerCfg	Policy optimizer configuration.

Functions:

bind_scheduler_horizon(cfg, total_updates)	Fill total_iters / T_max from the training update budget.
build_lr_scheduler(optimizer, cfg)	Build a scheduler, or None when name is unset.
build_optimizer(parameters[, cfg])
coerce_lr_scheduler_cfg(value)
coerce_optimizer_cfg(value)
dict_to_tensordict(obs_dict, device)	Convert an environment observation mapping into a TensorDict.
flatten_dict_observation(obs)	Flatten a hierarchical observation TensorDict into a 2D tensor.
get_registered_lr_scheduler_names()
get_registered_optimizer_names()
scheduler_needs_horizon(cfg)

class embodichain.learning.rl.utils.AlgorithmCfg

Bases: object

Shared fields for RL algorithm configs.

Methods:

__init__([device, optimizer, lr_scheduler, ...])
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:

batch_size
device
gae_lambda
gamma
lr_scheduler
max_grad_norm
optimizer

__init__(device=<factory>, optimizer=<factory>, lr_scheduler=<factory>, batch_size=<factory>, gamma=<factory>, gae_lambda=<factory>, max_grad_norm=<factory>)

batch_size: int

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.

device: str

gae_lambda: float

gamma: float

lr_scheduler: LRSchedulerCfg

max_grad_norm: float

optimizer: OptimizerCfg

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.

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.

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.utils.LRSchedulerCfg

Bases: object

Optional LR scheduler. name=None disables scheduling.

Horizon keys (total_iters / T_max) may be omitted and bound later by BaseAlgorithm.bind_schedule.

Methods:

__init__([name, kwargs])
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:

kwargs
name

__init__(name=<factory>, kwargs=<factory>)

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.

kwargs: dict[str, Any]

name: str | None

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.

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.

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.utils.OptimizerCfg

Bases: object

Policy optimizer configuration.

Methods:

__init__([name, learning_rate, kwargs])
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:

kwargs
learning_rate
name

__init__(name=<factory>, learning_rate=<factory>, kwargs=<factory>)

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.

kwargs: dict[str, Any]

learning_rate: float

name: str

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.

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.

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.

embodichain.learning.rl.utils.bind_scheduler_horizon(cfg, total_updates)

Fill total_iters / T_max from the training update budget.

Return type:

LRSchedulerCfg

embodichain.learning.rl.utils.build_lr_scheduler(optimizer, cfg)

Build a scheduler, or None when name is unset.

Return type:

LRScheduler | None

embodichain.learning.rl.utils.build_optimizer(parameters, cfg=None)

Return type:

Optimizer

embodichain.learning.rl.utils.coerce_lr_scheduler_cfg(value)

Return type:

LRSchedulerCfg

embodichain.learning.rl.utils.coerce_optimizer_cfg(value)

Return type:

OptimizerCfg

embodichain.learning.rl.utils.dict_to_tensordict(obs_dict, device)

Convert an environment observation mapping into a TensorDict.

Parameters:

• obs_dict (Tensor | TensorDict | Mapping[str, Any]) – Tensor or mapping returned by reset() or step().

• device (device | str) – Target device for the resulting TensorDict.

Return type:

TensorDict

Returns:

Observation TensorDict moved onto the target device.

embodichain.learning.rl.utils.flatten_dict_observation(obs)

Flatten a hierarchical observation TensorDict into a 2D tensor.

Parameters:

obs (TensorDict) – Observation TensorDict with batch dimension [num_envs].

Return type:

Tensor

Returns:

Flattened observation tensor of shape [num_envs, obs_dim].

embodichain.learning.rl.utils.get_registered_lr_scheduler_names()

Return type:

list[str]

embodichain.learning.rl.utils.get_registered_optimizer_names()

Return type:

list[str]

embodichain.learning.rl.utils.scheduler_needs_horizon(cfg)

Return type:

bool