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].
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, *)
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:

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