embodichain.learning.rl

On-policy reinforcement learning pipeline.

Algorithms (PPO/GRPO), rollout buffers, collectors, 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.
policies	RL policies and registration, loaded on demand.
policy_evaluation	Policy evaluation configuration for embodichain eval-policy.
train
utils	RL helper utilities: algorithm config, optimizers, and observation helpers.

Top-level APIs

DifferentiableTrainer	Coordinate APG updates, full rollouts, and optional TBPTT boundaries.
DifferentiableTrainerCfg	Configuration for segmented or complete graph-preserving training.
DifferentiableRolloutSpec	Describe one independent complete rollout used for a gradient microbatch.
DifferentiableVecEnv	Batched env that preserves the autograd path through step.
LearningVecEnv	Structural interface shared by lightweight vector environments.
ScheduledDifferentiableVecEnv	Differentiable env that schedules variable independent rollouts.
RunningObservationNormalizer	Normalize continuous observation fields with running Welford statistics.
BatchedGradientNormStats	Accumulate row-wise adjoint norm and clipping statistics on-device.
build_learning_env(name, *, num_envs, ...)	Build a registered lightweight vector environment.
clip_batched_gradient_norm(gradient, max_norm)	Clip each batch row without shortening the differentiable time horizon.
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.
stratified_rollout_value(index, minimum, maximum)	Cycle uniformly through an integer range and rotate each cycle's order.

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, *)
complete_discounted_return(rollout, gamma)	Compute per-environment returns up to the first terminal transition.
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_complete_rollout(rollout, *[, ...])	Accumulate one independent full-horizon rollout.
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_complete_rollout(rollout, *, objective_scale=1.0, accumulation_scale=1.0)

Accumulate one independent full-horizon rollout.

Parameters:

• rollout (DifferentiableRollout) – Complete graph-preserving trajectory.

• objective_scale (float | Tensor) – Scalar or per-environment return multiplier.

• accumulation_scale (float) – Loss multiplier, normally reciprocal to the number of independent rollout microbatches in one update.

Raises:

• RuntimeError – If no APG update is active.

• ValueError – If the rollout, scaling, or accumulation factor is invalid.

Return type:

None

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
max_grad_norm_before_clip
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>, max_grad_norm_before_clip=<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

max_grad_norm_before_clip: 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()
load_state_dict(state, *[, resume_scheduler])	Restore optimization state or restart its configured LR curve.
state_dict()	Return optimizer and learning-rate state for a trainer checkpoint.
update(rollout)	Update policy using collected data and return training losses.
validate_state_dict(state, *[, resume_scheduler])	Check optimization structure before changing the live algorithm.

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

load_state_dict(state, *, resume_scheduler='restore')

Restore optimization state or restart its configured LR curve.

Parameters:

• state (dict[str, Any]) – Optimization fields returned by state_dict().

• resume_scheduler (str) – Restore the saved LR curve, or restart the configured curve while retaining optimizer momentum.

Return type:

None

lr_scheduler: LRScheduler | None

optimizer: Optimizer

rollout_kind = 'standard'

state_dict()

Return optimizer and learning-rate state for a trainer checkpoint.

Return type:

dict[str, Any]

Returns:

Optimizer state, scheduler configuration/progress, and any pending restart that binds the schedule to the remaining update budget.

abstract update(rollout)

Update policy using collected data and return training losses.

Return type:

Dict[str, float]

validate_state_dict(state, *, resume_scheduler='restore')

Check optimization structure before changing the live algorithm.

Parameters:

• state (dict[str, Any]) – Optimization fields returned by state_dict().

• resume_scheduler (str) – Restore the saved curve or restart configuration.

Raises:

ValueError – If the mode, optimizer layout, or scheduler fields differ.

Return type:

None

class embodichain.learning.rl.algo.GRPO

Bases: BaseAlgorithm[TensorDict]

Group Relative Policy Optimization on top of TensorDict rollouts.

Methods:

__init__(cfg, policy)
state_dict()	Include the fixed KL reference policy in optimization state.
update(rollout)	Update policy using collected data and return training losses.
validate_state_dict(state, *[, resume_scheduler])	Reject checkpoints missing the active KL reference policy.

__init__(cfg, policy)

state_dict()

Include the fixed KL reference policy in optimization state.

Return type:

dict[str, Any]

Returns:

Optimization state with the KL reference when enabled.

update(rollout)

Update policy using collected data and return training losses.

Return type:

Dict[str, float]

validate_state_dict(state, *, resume_scheduler='restore')

Reject checkpoints missing the active KL reference policy.

Parameters:

• state (dict[str, Any]) – Optimization and reference-policy checkpoint fields.

• resume_scheduler (str) – Restore the saved curve or restart configuration.

Raises:

ValueError – If the optimizer or active reference policy is incompatible.

Return type:

None

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:

action_bound_coef
batch_size
clip_coef
desired_kl
device
ent_coef
gae_lambda
gamma
lr_scheduler
max_grad_norm
maximum_learning_rate
minimum_learning_rate
n_epochs
normalize_advantage_per_mini_batch
num_mini_batches
optimizer
schedule
use_clipped_value_loss
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>, action_bound_coef=<factory>, use_clipped_value_loss=<factory>, schedule=<factory>, desired_kl=<factory>, minimum_learning_rate=<factory>, maximum_learning_rate=<factory>, num_mini_batches=<factory>, normalize_advantage_per_mini_batch=<factory>)

action_bound_coef: float

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.

desired_kl: float | None

device: str

ent_coef: float

gae_lambda: float

gamma: float

lr_scheduler: LRSchedulerCfg

max_grad_norm: float

maximum_learning_rate: float

minimum_learning_rate: float

n_epochs: int

normalize_advantage_per_mini_batch: bool

num_mini_batches: int | None

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.

schedule: str

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.

use_clipped_value_loss: 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.

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.complete_discounted_return(rollout, gamma)

Compute per-environment returns up to the first terminal transition.

Unlike segmented TBPTT returns, complete-rollout discounting never restarts after done. This excludes rewards emitted by any automatic reset during the fixed full horizon.

Parameters:

• rollout (DifferentiableRollout) – Complete graph-preserving trajectory.

• gamma (float) – Per-step discount factor.

Return type:

Tensor

Returns:

One masked discounted return per environment.

Raises:

ValueError – If the rollout is empty.

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:

DifferentiableRolloutSpec	Describe one independent complete rollout used for a gradient microbatch.
DifferentiableVecEnv	Batched env that preserves the autograd path through step.
LearningVecEnv	Structural interface shared by lightweight vector environments.
ScheduledDifferentiableVecEnv	Differentiable env that schedules variable independent rollouts.

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.
stratified_rollout_value(index, minimum, maximum)	Cycle uniformly through an integer range and rotate each cycle's order.

class embodichain.learning.rl.env.DifferentiableRolloutSpec

Bases: object

Describe one independent complete rollout used for a gradient microbatch.

Parameters:

• num_steps (int) – Full rollout horizon. The trainer must not detach or truncate it.

• objective_scale (float | Tensor) – Per-environment or scalar multiplier applied to returns.

• metadata (Mapping[str, float]) – Scalar labels recorded with training metrics.

Methods:

__init__(num_steps[, objective_scale, metadata])

Attributes:

metadata
num_steps
objective_scale

__init__(num_steps, objective_scale=1.0, metadata=<factory>)

metadata: Mapping[str, float]

num_steps: int

objective_scale: float | Tensor = 1.0

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]]

class embodichain.learning.rl.env.ScheduledDifferentiableVecEnv

Bases: DifferentiableVecEnv, Protocol

Differentiable env that schedules variable independent rollouts.

The trainer calls prepare_differentiable_rollout() immediately before resetting the environment. Implementations may select task difficulty for the next reset, such as an ordered-waypoint count, and must return the full horizon and objective scaling for that selection.

Methods:

prepare_differentiable_rollout(rollout_index)

Configure the next reset and return its complete-rollout contract.

prepare_differentiable_rollout(rollout_index)

Configure the next reset and return its complete-rollout contract.

Parameters:

rollout_index (int) – Zero-based independent-rollout index.

Return type:

DifferentiableRolloutSpec

Returns:

Full horizon, objective scaling, and optional metric metadata.

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]

embodichain.learning.rl.env.stratified_rollout_value(index, minimum, maximum)

Cycle uniformly through an integer range and rotate each cycle’s order.

Parameters:

• index (int) – Zero-based rollout index.

• minimum (int) – Inclusive minimum scheduled value.

• maximum (int) – Inclusive maximum scheduled value.

Return type:

int

Returns:

Scheduled integer for index.

Raises:

ValueError – If the inclusive range is empty.

Evaluation

Shared deterministic episode evaluation for all RL trainers.

Functions:

convert_policy_action_for_env(env, action)	Move a flat policy action to the task environment's device.
evaluate_episodes(*, policy, env, ...[, ...])	Evaluate exactly num_episodes completed asynchronous episodes.
infer_policy_action(policy, observation, *, ...)	Run the same deterministic Policy call used by RL evaluation.
prepare_policy_observation(observation, device)	Flatten one Environment observation in the configured input order.

embodichain.learning.rl.evaluation.convert_policy_action_for_env(env, action)

Move a flat policy action to the task environment’s device.

Parameters:

• env (Any) – Environment that owns action preprocessing and simulation state.

• action (Tensor) – Flat action produced on the Policy inference device.

Return type:

Any

Returns:

Action on the Environment device. The environment owns action-term processing so that policy actions are validated and applied exactly once inside EmbodiedEnv.step().

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

Evaluate exactly num_episodes completed asynchronous episodes.

Parameters:

• policy (Module) – Policy evaluated with deterministic actions.

• env (Any) – Vector environment with automatic row reset.

• num_episodes (int) – Exact number of completed episodes to collect.

• device (device | str) – Policy device.

• seed (int | None) – Optional environment reset seed.

• on_step (Callable[[dict[str, Any]], None] | None) – Optional callback for raw step info.

• observation_transform (Callable[[Tensor], Tensor] | None) – Optional frozen preprocessing transform, such as a running observation normalizer.

Return type:

dict[str, float]

embodichain.learning.rl.evaluation.infer_policy_action(policy, observation, *, device, num_envs, observation_transform=None)

Run the same deterministic Policy call used by RL evaluation.

Return type:

Tensor

embodichain.learning.rl.evaluation.prepare_policy_observation(observation, device, groups=None)

Flatten one Environment observation in the configured input order.

Return type:

Tensor

Gradient Stabilization

Gradient stabilization primitives for differentiable rollouts.

Classes:

BatchedGradientNormStats

Accumulate row-wise adjoint norm and clipping statistics on-device.

Functions:

clip_batched_gradient_norm(gradient, max_norm)

Clip each batch row without shortening the differentiable time horizon.

class embodichain.learning.rl.gradients.BatchedGradientNormStats

Bases: object

Accumulate row-wise adjoint norm and clipping statistics on-device.

Parameters:

device (device | str) – Device on which hook-side counters are accumulated.

Methods:

__init__(device)
metrics()	Return statistics after backward has invoked registered hooks.

__init__(device)

metrics()

Return statistics after backward has invoked registered hooks.

Return type:

dict[str, float]

Returns:

Mean and maximum pre-clip norm plus clipped/non-finite fractions.

embodichain.learning.rl.gradients.clip_batched_gradient_norm(gradient, max_norm, stats=None)

Clip each batch row without shortening the differentiable time horizon.

Norms are computed with max-absolute-value scaling to avoid overflow in float32. A non-finite row is replaced with zeros while finite rows remain independent from one another.

Parameters:

• gradient (Tensor) – Tensor whose first dimension identifies independent rows.

• max_norm (float) – Maximum L2 norm per row. Zero disables clipping.

• stats (BatchedGradientNormStats | None) – Optional on-device accumulator populated before clipping.

Return type:

Tensor

Returns:

A finite tensor with the same shape, dtype, and device as gradient.

Raises:

ValueError – If max_norm is negative or gradient is not batched.

Observation Normalization

Running statistics for learning-environment observations.

Classes:

RunningObservationNormalizer

Normalize continuous observation fields with running Welford statistics.

class embodichain.learning.rl.normalization.RunningObservationNormalizer

Bases: object

Normalize continuous observation fields with running Welford statistics.

The optional mask leaves semantic fields such as one-hot encodings and validity bits unchanged while still tracking a single flat observation. Statistics are updated only when update() is called; normalize() is side-effect free so a complete differentiable rollout uses one frozen normalization transform.

Parameters:

• observation_dim (int) – Flat observation dimension.

• device (device | str) – Device used for the running statistics.

• normalize_mask (Tensor | None) – Boolean mask where True selects normalized fields.

• initial_count (float) – Positive pseudo-count used to stabilize the first update.

Methods:

__init__(observation_dim, device[, ...])
load_state_dict(state_dict)	Restore statistics while validating their observation layout.
normalize(observations)	Apply the frozen running transform without detaching observations.
state_dict()	Return a device-independent checkpoint payload.
update(observations)	Merge one observation batch into the running statistics.

__init__(observation_dim, device, normalize_mask=None, *, initial_count=0.0001)

load_state_dict(state_dict)

Restore statistics while validating their observation layout.

Parameters:

state_dict (Mapping[str, Any]) – Payload produced by state_dict().

Raises:

ValueError – If dimensions or the pseudo-count are invalid.

Return type:

None

normalize(observations)

Apply the frozen running transform without detaching observations.

Parameters:

observations (Tensor) – Tensor ending in the configured observation dimension.

Return type:

Tensor

Returns:

Tensor with continuous fields normalized and semantic fields intact.

Raises:

ValueError – If the trailing observation dimension is incompatible.

state_dict()

Return a device-independent checkpoint payload.

Return type:

dict[str, Any]

Returns:

Mapping containing mean, variance, count, and normalization mask.

update(observations)

Merge one observation batch into the running statistics.

Parameters:

observations (Tensor) – Finite tensor shaped [batch, observation_dim].

Raises:

ValueError – If the shape is incompatible or values are non-finite.

Return type:

None

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, full rollouts, and optional TBPTT boundaries.
DifferentiableTrainerCfg	Configuration for segmented or complete graph-preserving training.

class embodichain.learning.rl.differentiable_trainer.DifferentiableTrainer

Bases: object

Coordinate APG updates, full rollouts, and optional TBPTT boundaries.

Methods:

__init__(cfg, env, policy, algorithm[, ...])
get_summary()	Return the current in-memory training summary.
load_checkpoint(path, *[, resume_scheduler])	Validate and restore learning state, starting with a fresh rollout.
save_checkpoint([path])	Atomically save learning state and update the run index.
train([total_timesteps, total_updates])	Train to a vector-transition or optimizer-update budget.

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

get_summary()

Return the current in-memory training summary.

Return type:

dict[str, Any]

load_checkpoint(path, *, resume_scheduler='restore')

Validate and restore learning state, starting with a fresh rollout.

Parameters:

• path (str | Path) – Checkpoint file, including existing APG schema-v1 files.

• resume_scheduler (str) – Restore the saved LR curve or restart the configured curve while preserving optimizer momentum.

Return type:

None

save_checkpoint(path=None)

Atomically save learning state and update the run index.

Parameters:

path (str | Path | None) – Destination, or the current step’s default checkpoint path.

Return type:

str

Returns:

Saved checkpoint path.

train(total_timesteps=None, *, total_updates=None)

Train to a vector-transition or optimizer-update budget.

total_updates is the stable budget for scheduled variable-horizon training. Exactly one budget may be supplied.

Parameters:

• total_timesteps (int | None) – Optional absolute vector-transition budget.

• total_updates (int | None) – Optional absolute optimizer-update budget.

Return type:

dict[str, Any]

Returns:

Current training counters, metrics, histories, and checkpoint paths.

Raises:

ValueError – If neither/both budgets are supplied or one is negative.

class embodichain.learning.rl.differentiable_trainer.DifferentiableTrainerCfg

Bases: object

Configuration for segmented or complete graph-preserving 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:

action_adjoint_max_norm
best_eval_metric
best_eval_mode
checkpoint_dir
clip_actions_to_space
deterministic_actions
eval_frequency_steps
eval_seed
experiment_name
gradient_accumulation_steps
normalize_observations
num_eval_episodes
rollout_mode
rollout_seed
save_frequency_updates
segment_length
update_horizon
use_wandb

__init__(segment_length=<factory>, update_horizon=<factory>, rollout_mode=<factory>, gradient_accumulation_steps=<factory>, deterministic_actions=<factory>, clip_actions_to_space=<factory>, action_adjoint_max_norm=<factory>, normalize_observations=<factory>, rollout_seed=<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>)

action_adjoint_max_norm: float

best_eval_metric: str

best_eval_mode: str

checkpoint_dir: str

clip_actions_to_space: bool

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

gradient_accumulation_steps: int

normalize_observations: bool

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.

rollout_mode: str

rollout_seed: int | None

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, critic_obs_dim=None, distribution_param_dim=None)

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, *, observation_normalizer=None, clip_actions_to_space=False, action_adjoint_max_norm=0.0)

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:

action_gradient_stats
alive_mask	Return rows active before each step, stopping after the first done.
final_observation	Return the observation after the final transition.
initial_observation
num_steps	Return the number of collected transitions.
observations	Stack the raw policy observations.
rewards	Stack rewards as [time, num_envs] without detaching them.
transitions

__init__(initial_observation, transitions, action_gradient_stats=None)

action_gradient_stats: BatchedGradientNormStats | None = None

property alive_mask: Tensor

Return rows active before each step, stopping after the first done.

Returns:

Boolean tensor shaped [time, num_envs].

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 observations: Tensor

Stack the raw policy observations.

Returns:

Tensor shaped [time, num_envs, features].

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

Policies

Policy implementations and registration live in embodichain.learning.rl.policies.

• embodichain.learning.rl.policies

Runtime Reconstruction

Shared environment and Policy construction for RL training and evaluation.

Classes:

PolicyRuntime

An Environment and Policy reconstructed from one training configuration.

Functions:

build_gym_policy_runtime(config, *, device, ...)	Build a simulator task and the Policy declared by its training config.
build_learning_policy_runtime(config, *, ...)	Build a lightweight Environment and its configured Policy.

class embodichain.learning.rl.runtime.PolicyRuntime

Bases: object

An Environment and Policy reconstructed from one training configuration.

Methods:

__init__(env, policy, device, env_id[, ...])
close()	Close the Environment without terminating the current process.

Attributes:

device
env
env_cfg
env_id
gym_config
gym_config_path
policy

__init__(env, policy, device, env_id, env_cfg=None, gym_config=None, gym_config_path=None)

close()

Close the Environment without terminating the current process.

Return type:

None

device: device

env: Any

env_cfg: Any | None = None

env_id: str

gym_config: dict[str, Any] | None = None

gym_config_path: Path | None = None

policy: Module

embodichain.learning.rl.runtime.build_gym_policy_runtime(config, *, device, num_envs, headless, renderer, gpu_id, seed=None, config_dir=None, profiler=None, simulation_device=None, trainer_events=())

Build a simulator task and the Policy declared by its training config.

Parameters:

• config (dict[str, Any]) – Training configuration containing trainer and policy blocks.

• device (device) – Policy tensor device.

• num_envs (int | None) – Optional override for the number of environments.

• headless (bool) – Whether to suppress the native window.

• renderer (str | None) – Explicit renderer override; None preserves the task’s render configuration, including demand-based auto selection.

• gpu_id (int) – Renderer GPU index when the simulation device has no CUDA index.

• seed (int | None) – Optional environment seed override.

• config_dir (str | Path | None) – Directory used to resolve the referenced task configuration.

• profiler (ProfilerCfg | None) – Optional environment profiler configuration.

• simulation_device (device | None) – Optional simulation device distinct from the policy.

• trainer_events (Sequence[EventCfg]) – Active training and evaluation events, resolved before the first World so their cameras can reserve the process renderer.

Return type:

PolicyRuntime

Returns:

The constructed environment, policy and resolved task configuration.

embodichain.learning.rl.runtime.build_learning_policy_runtime(config, *, device, num_envs)

Build a lightweight Environment and its configured Policy.

Return type:

PolicyRuntime

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_task 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, resume=None, run=None, resume_scheduler='restore')

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.

• resume (str | Path | None) – Checkpoint file, or latest/best selected from run.

• run (str | Path | None) – Source run directory when resume selects latest/best.

• resume_scheduler (str) – Restore the saved LR curve or restart the configured curve with optimizer momentum retained.

Return type:

dict[str, object] | None

Returns:

The lightweight trainer summary, or None for simulator training.

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.
flatten_observation_groups(obs, groups)	Flatten selected observation groups in the configured order.
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.flatten_observation_groups(obs, groups)

Flatten selected observation groups in the configured order.

A group name selects one top-level TensorDict key. Dot-separated names may select nested keys, for example "robot.qpos". When groups is None, every tensor is flattened using flatten_dict_observation() for compatibility with existing tasks.

Parameters:

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

• groups (Sequence[str] | None) – Observation group names consumed by one model.

Return type:

Tensor

Returns:

Concatenated observation tensor with shape [num_envs, obs_dim].

Raises:

• ValueError – If groups is empty or a selected group has no tensors.

• KeyError – If a configured group is absent from obs.

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