embodichain.learning.rl.algo

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

Overview

Algorithm registry and algorithm-construction helpers for RL training. The on-policy algorithms PPO and GRPO and the differentiable algorithm APG all derive from BaseAlgorithm. Each algorithm declares a RolloutKind that selects the compatible trainer. build_algo() looks up a registered algorithm by name and wires it to a policy, while compute_gae() provides generalized advantage estimation.

Classes

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

Functions

build_algo(name, cfg_kwargs, policy, device, *)
get_registered_algo_names()
compute_gae(rollout, gamma, gae_lambda)	Compute GAE over a rollout stored as [num_envs, time + 1].
complete_discounted_return(rollout, gamma)	Compute per-environment returns up to the first terminal transition.
segmented_discounted_return(rollout, gamma)	Compute one discounted return per environment within a rollout segment.

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:

device
lr_scheduler
optimizer
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