Add GPU differentiable IK layers and shared neural networks

[yuecideng]

Description

Add GPU-batched differentiable kinematics and fixed-step IK for training and task-parameter optimization. Target poses and initial joint states retain first-order gradients through FK, the Jacobian and every DLS update. Failed solves return their final iterate and continuous residuals.

The ownership boundaries are:

• compute/kinematics: pure Torch DLS, pose residuals, copied PK geometry, configuration and result types.
• learning/layers: KinematicsLayer and DifferentiableIKLayer, with float32/float64, device migration and fixed-geometry checkpoint validation. Forward-only inference uses the same layer under torch.no_grad().
• learning/models: only general components — MLP, EmpiricalNormalizer and build_mlp_from_cfg. Shared networks, feature normalization and MLP construction do not import RL, TensorDict or simulation.
• learning/rl/policies: RL-specific Policy, ActorOnly, ActorCritic, policy registration, observation-group resolution and policy construction. RL runtime, collectors, algorithms and trainers consume these interfaces and the shared components.

Generic components import from embodichain.learning.models; RL abstractions import from embodichain.learning.rl.policies. There are no compatibility aliases between these namespaces. External imports must use learning.models for generic components and learning.rl.policies for RL policies. The previous learning.rl.models package is removed. RL policy behavior and state-dict structure are preserved.

Assigned checkpoint loads synchronize fixed PK geometry to the loaded tensor buffers through a child-module post-load hook. This also works when an enclosing network loads its state dict. Synchronization is confined to checkpoint loading; FK and IK calls retain their existing computation path.

Callers own root-frame conversion, TCP, physical limits and optimization. Hard-limit projection and the rotation branch cut have documented gradient boundaries. Robot defaults and registered runtime solvers are unchanged. The unified benchmark measures differentiable forward/backward and component costs. The runnable use case is the interactive grasp-yaw demo in #808; the earlier synthetic training example is removed.

API references, learning/RL guides, context ownership and architecture evidence reflect these boundaries. Dependencies: existing project packages; no new requirements. No linked issue.

Stack

• Layer: 1/2
• Base: main
• Depends on: none
• Followed by: Add a learned grasp-yaw demo with differentiable IK #808

Type of change

• New feature
• Breaking change: generic network/normalizer/builder imports move to learning.models; RL policy imports move to learning.rl.policies.

Validation

All current validation uses conda py311:

• Black 26.3.1: python -m black .; git diff --check.
• API documentation coverage: 2619/2619 exports documented; context checker and ownership routes passed.
• Shared components, standalone layers, APG, differentiable collector/trainer and context regressions: 81 passed, 2 CUDA tests skipped in the CPU invocation.
• Shared-component tests verify a configured network's supervised gradients/initialization and feature statistics, checkpoint restoration and input gradients. A blocked-import subprocess verifies that generic MLP construction and normalization avoid RL, TensorDict and simulation; the policy registry and lightweight APG imports are checked separately, including the new public module identity and failure to import the removed legacy Policy API even when stale bytecode directories remain.
• PPO stability, runtime/checkpoint reconstruction, routing and point-mass regressions: 49 passed using the process-local DexSim import overlay described below.
• Dedicated CUDA float32/float64 parity and synchronization checks: 2 passed before this package-only rename; unchanged numerical CUDA behavior was not rebenchmarked.
• API-checker and architecture regression tests: 55 passed. Architecture generation and snapshot validation passed.
• Sphinx dummy build completed with the same process-local overlay; model API pages have no remaining warnings. Pre-existing repository-wide warnings and unrelated docstring parsing diagnostics remain, so this is not a warnings-as-errors build claim.
• Checkpoint-loading regression fix: 69 CPU tests passed, 4 CUDA tests skipped in that invocation; a separate serial CUDA run passed 4 tests, including float32/float64 parity and assigned loads in both CPU↔CUDA directions.
• Before the fix, all six new cross-precision checkpoint cases failed at FK/Jacobian evaluation. They now pass for standalone FK, IK, and enclosing-network loads; restored IK residuals and differentiable FK/Jacobian outputs match the source. Unsupported assigned float16 buffers are rejected.
• The dependent grasp-yaw demo is revalidated on this base in Add a learned grasp-yaw demo with differentiable IK #808.

The installed DexSim generated exports lack cable-render kernel exports. The 49 runtime regressions use an existing process-local overlay exposing the actual cable kernels from the sibling DexSim sources; no dependency sources or test stubs are changed. Standalone checks and the grasp-yaw demo run without this overlay.

Context review: revise learning-layers and RL ownership, retain policy registration in RL, synchronize model API/architecture references, and document assigned-checkpoint geometry migration. The checkpoint fix preserves architecture topology and public exports; its existing architecture snapshot remains valid.

Screenshots

Interactive preview and demo artifacts are in #808.

Checklist

• Run Black 26.3.1 over the repository.
• Update affected documentation and agent context.
• Update public API references and pass coverage checks.
• Run focused numerical, gradient, shared-component and RL runtime regressions.
• Review dependencies; no new package requirements.

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge with the documented import migration.

Summary

Adds standalone Torch kinematics and fixed-step differentiable IK, alongside shared neural-network components.

• IK layers return joint results while keeping gradients connected.
• Shared networks and RL policies now live in separate namespaces.
• The kinematics benchmark now measures IK training work.
• Gym setup passes an isolated robot config into simulation.
• Architecture records now reflect the updated repository layout.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  C["compute.kinematics: FK, Jacobians and DLS"] --> L["learning.layers: Torch modules"]
  M["learning.models: MLP and normalization"] --> P["learning.rl.policies: policies and registry"]
  P --> R["RL training and evaluation"]
  L --> O["Caller-owned losses and optimization"]

Loading

_{Reviews (6) · Last reviewed commit: "fix(learning): synchronize kinematics af..." · Reviewed by Greptile}