embodichain.learning.rl.policies

RL policy interfaces and construction

Policy defines the action, action-evaluation and value interfaces used by RL algorithms. ActorCritic combines a Gaussian action policy with a value network; ActorOnly supplies the same rollout contract without a learned value function. The registry builds policies from registered names and observation-group settings. General MLP construction and feature normalization are provided separately by embodichain.learning.models.

Policy	Abstract base class that all RL policies must implement.
ActorCritic	TensorDict actor-critic with a learnable Gaussian action distribution.
ActorOnly	Actor-only policy for algorithms that do not use a value function (e.g., GRPO).
register_policy(name, policy_cls)
get_registered_policy_names()
get_policy_class(name)
build_policy(policy_block, obs_space, ...[, ...])	Build a policy from config using spaces for extensibility.
resolve_policy_obs_groups(policy_block)	Resolve actor and critic observation groups from policy configuration.

RL policies and registration, loaded on demand.

Classes:

ActorCritic	TensorDict actor-critic with a learnable Gaussian action distribution.
ActorOnly	Actor-only policy for algorithms that do not use a value function (e.g., GRPO).
Policy	Abstract base class that all RL policies must implement.

Functions:

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)
resolve_policy_obs_groups(policy_block)	Resolve actor and critic observation groups from policy configuration.

class embodichain.learning.rl.policies.ActorCritic

Bases: Policy

TensorDict actor-critic with a learnable Gaussian action distribution.

Parameters:

• obs_dim (int) – Actor observation dimension.

• action_dim (int) – Action dimension.

• device (device) – Device used by the policy modules and distribution parameters.

• actor (Module) – Module that maps actor observations to action means.

• critic (Module) – Module that maps critic observations to scalar values.

• critic_obs_dim (int | None) – Separate critic observation dimension. None reuses actor observations.

• actor_obs_groups (Sequence[str] | None) – Ordered observation groups consumed by the actor.

• critic_obs_groups (Sequence[str] | None) – Ordered observation groups consumed by the critic.

• actor_obs_normalization (bool) – Whether to normalize actor observations.

• critic_obs_normalization (bool) – Whether to normalize critic observations.

• initial_action_std (float) – Initial standard deviation for every action.

• action_std_range (tuple[float, float]) – Inclusive clamp range applied to standard deviation.

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.
optimization_parameter_groups()	Return actor-side and critic-side trainable parameter groups.
update_normalization(tensordict)	Update actor and critic observation statistics from environment data.

Attributes:

action_std

Return the clamped state-independent action standard deviation.

__init__(obs_dim, action_dim, device, actor, critic, critic_obs_dim=None, actor_obs_groups=None, critic_obs_groups=None, actor_obs_normalization=False, critic_obs_normalization=False, initial_action_std=1.0, action_std_range=(1e-06, 1000000.0))

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

property action_std: Tensor

Return the clamped state-independent action standard deviation.

evaluate_actions(tensordict)

Evaluate actions and return current policy outputs.

Parameters:

tensordict (TensorDict) – TensorDict containing obs, action, and optional critic_obs.

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 and, for policies with a privileged critic, critic_obs.

Return type:

TensorDict

Returns:

TensorDict with value populated.

optimization_parameter_groups()

Return actor-side and critic-side trainable parameter groups.

Return type:

tuple[tuple[Parameter, ...], tuple[Parameter, ...]]

update_normalization(tensordict)

Update actor and critic observation statistics from environment data.

Return type:

None

class embodichain.learning.rl.policies.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, actor_obs_groups=None)

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, action, and optional critic_obs.

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 and, for policies with a privileged critic, critic_obs.

Return type:

TensorDict

Returns:

TensorDict with value populated.

class embodichain.learning.rl.policies.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, action, and optional critic_obs.

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 and, for policies with a privileged critic, critic_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 and, for policies with a privileged critic, critic_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 and, for policies with a privileged critic, critic_obs.

Return type:

TensorDict

Returns:

TensorDict with value populated.

training: bool

embodichain.learning.rl.policies.build_policy(policy_block, obs_space, action_space, device, actor=None, critic=None, critic_obs_space=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.policies.get_policy_class(name)

Return type:

Optional[Type[Policy]]

embodichain.learning.rl.policies.get_registered_policy_names()

Return type:

list[str]

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

Return type:

None

embodichain.learning.rl.policies.resolve_policy_obs_groups(policy_block)

Resolve actor and critic observation groups from policy configuration.

Parameters:

policy_block (Mapping[str, object]) – Policy configuration containing an optional obs_groups mapping.

Return type:

tuple[tuple[str, ...] | None, tuple[str, ...] | None]

Returns:

Actor and critic observation group names. Both are None when the existing flatten-all behavior is selected.

Raises:

• TypeError – If the observation-group configuration is malformed.

• ValueError – If the actor observation set is empty.

Policy registry

RL policy registration and construction for actor-critic and actor-only models.

Functions:

build_policy(policy_block, obs_space, ...[, ...])	Build a policy from config using spaces for extensibility.
resolve_policy_obs_groups(policy_block)	Resolve actor and critic observation groups from policy configuration.

embodichain.learning.rl.policies.registry.build_policy(policy_block, obs_space, action_space, device, actor=None, critic=None, critic_obs_space=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.policies.registry.resolve_policy_obs_groups(policy_block)

Resolve actor and critic observation groups from policy configuration.

Parameters:

policy_block (Mapping[str, object]) – Policy configuration containing an optional obs_groups mapping.

Return type:

tuple[tuple[str, ...] | None, tuple[str, ...] | None]

Returns:

Actor and critic observation group names. Both are None when the existing flatten-all behavior is selected.

Raises:

• TypeError – If the observation-group configuration is malformed.

• ValueError – If the actor observation set is empty.

Observation normalization

ActorCritic can maintain independent actor and critic observation moments. It uses the shared EmpiricalNormalizer; policy buffers preserve the statistics in checkpoints. Observation selection and normalization-update timing remain part of the RL rollout contract.