embodichain.learning

Learning layers, shared models and reinforcement-learning systems.

Subpackages load on demand so standalone layers do not import RL or simulation.

Submodules

layers	Simulation-independent differentiable robotics layers for Torch networks.
models	General neural networks, normalization and network construction.
rl	On-policy reinforcement learning pipeline.
sl	Supervised learning (SL) for policies and learned state transitions.

Reinforcement Learning

On-policy reinforcement learning pipeline.

Algorithms (PPO/GRPO), rollout buffers, collectors, and the training entry point; rollout data flows as TensorDict objects.

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.
train
utils	RL helper utilities: algorithm config, optimizers, and observation helpers.

Differentiable layers

learning.layers composes simulation-independent computation as Torch modules. KinematicsLayer owns fixed TCP and limit buffers; its copied PK geometry follows to(device/dtype). Reconstruct the same robot model before loading a checkpoint; its geometry fingerprint is checked. DifferentiableIKLayer.forward returns joints for network composition, while solve returns residuals and a success mask. Both retain input gradients in training and evaluation mode. Callers own rollout state, optimizers, and detach boundaries.

KinematicsLayer	Expose FK and differentiable Jacobians as a reusable network layer.
DifferentiableIKLayer	Compose fixed-step DLS with a differentiable kinematics module.

Torch module interface to fixed robot geometry, independent of simulation.

Classes:

KinematicsLayer

Expose FK and differentiable Jacobians as a reusable network layer.

class embodichain.learning.layers.kinematics.KinematicsLayer

Expose FK and differentiable Jacobians as a reusable network layer.

Parameters:

• model (SerialChainKinematics) – Fixed kinematics model, copied so device moves are independent.

• tcp (Tensor | None) – Optional flange-to-tool rigid transform.

• lower (Tensor | None) – Optional physical lower limits; defaults to the chain limits.

• upper (Tensor | None) – Optional physical upper limits; defaults to the chain limits.

Note

Geometry is configured, not trained. Checkpoints store TCP and limits and verify the identity of the model supplied at reconstruction time. load_state_dict(assign=True) adopts checkpoint device and precision for both tensor buffers and fixed geometry, including nested loads. train() and eval() do not disable input gradients.

Methods:

__init__(model[, tcp, lower, upper])	Initialize internal Module state, shared by both nn.Module and ScriptModule.
default_seed()	Return bounded joint midpoints, or projected zero for unbounded joints.
evaluate(qpos)	Return TCP transforms and root-frame geometric Jacobians.
forward(qpos)	Return root-frame TCP transforms for (..., dof) joints.
get_extra_state()	Record fixed geometry identity alongside tensor buffers.
jacobian(qpos)	Return differentiable TCP Jacobians, shape (..., 6, dof).
set_extra_state(state)	Reject checkpoints belonging to a different kinematic model.

Attributes:

dof	Number of movable joints in solver order.

__init__(model, tcp=None, lower=None, upper=None)

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

default_seed()

Return bounded joint midpoints, or projected zero for unbounded joints.

Return type:

Tensor

property dof: int

Number of movable joints in solver order.

evaluate(qpos)

Return TCP transforms and root-frame geometric Jacobians.

Return type:

tuple[Tensor, Tensor]

forward(qpos)

Return root-frame TCP transforms for (..., dof) joints.

Return type:

Tensor

get_extra_state()

Record fixed geometry identity alongside tensor buffers.

Return type:

dict[str, str]

jacobian(qpos)

Return differentiable TCP Jacobians, shape (..., 6, dof).

Return type:

Tensor

set_extra_state(state)

Reject checkpoints belonging to a different kinematic model.

Return type:

None

A graph-preserving IK layer for supervised and policy learning.

Classes:

DifferentiableIKLayer

Compose fixed-step DLS with a differentiable kinematics module.

class embodichain.learning.layers.ik.DifferentiableIKLayer

Compose fixed-step DLS with a differentiable kinematics module.

Parameters:

• kinematics (KinematicsLayer) – Model owning geometry, TCP and physical limit buffers.

• cfg (DifferentiableIKCfg | None) – Numerical settings, copied at construction.

Note

There is no recurrent seed state or optimizer inside this module. The caller owns rollout state and backpropagation boundaries.

Methods:

__init__(kinematics[, cfg])	Initialize internal Module state, shared by both nn.Module and ScriptModule.
forward(target_xpos[, qpos_seed])	Return joint positions while retaining target and seed gradients.
get_extra_state()	Save numerical settings with the module checkpoint.
set_extra_state(state)	Restore and validate numerical settings.
solve(target_xpos[, qpos_seed])	Return final joints, validity and differentiable TCP residuals.

__init__(kinematics, cfg=None)

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

forward(target_xpos, qpos_seed=None)

Return joint positions while retaining target and seed gradients.

Return type:

Tensor

get_extra_state()

Save numerical settings with the module checkpoint.

Return type:

dict[str, Any]

set_extra_state(state)

Restore and validate numerical settings.

Return type:

None

solve(target_xpos, qpos_seed=None)

Return final joints, validity and differentiable TCP residuals.

Parameters:

• target_xpos (Tensor) – Root-frame TCP transforms, shape (..., 4, 4).

• qpos_seed (Tensor | None) – (dof,) or matching batch dimensions; None uses midpoints.

Return type:

DifferentiableIKResult

Returns:

Result retaining the target’s leading batch dimensions.

Shared trainable models

learning.models owns general MLPs, feature normalization and MLP configuration builders. RL-specific Policy, ActorCritic, ActorOnly and policy registration/builders remain in learning.rl.policies. Shared components and standalone IK layers do not import RL, simulation or TensorDict. Concrete model families and their adapters are documented in the official model library.

• embodichain.learning.models

Supervised learning

Policies and learned dynamics share tensor-tree batches, objective-owned losses, optimizer updates and committed-update recovery through learning.sl.

• Supervised learning (SL)