Add a differentiable FK layer and pose-reaching demo

agent_context/MAP.yaml

@@ -291,9 +291,9 @@ topics:
   status: active
 - id: learning-layers
   title: Differentiable Learning Layers
-  aliases: [learning layers, differentiable IK layer, 可微 IK, 训练层, 学习模块分层]
-  keywords: [DifferentiableIKLayer, KinematicsLayer, SerialChainKinematics, DifferentiableIKCfg,
-    solve_dls_ik, damped_least_squares_step, Jacobian gradient, dJ/dq, learn_grasp_yaw, grasp yaw, 抓取角度优化, 抓取姿态学习]
+  aliases: [learning layers, differentiable FK layer, differentiable IK layer, 可微 FK, 可微 IK, 训练层, 学习模块分层]
+  keywords: [DifferentiableFKLayer, DifferentiableIKLayer, SerialChainKinematics, DifferentiableIKCfg,
+    solve_dls_ik, damped_least_squares_step, Jacobian gradient, dJ/dq, learn_grasp_yaw, learn_pose_reaching, grasp yaw, 抓取角度优化, 抓取姿态学习]
   paths: [topics/learning-layers/learning-layers.md]
   source_of_truth:
   - embodichain/compute/kinematics/dls.py
@@ -305,10 +305,13 @@ topics:
   - embodichain/learning/models/
   - examples/learning/differentiable_ik/learn_grasp_yaw.py
   - examples/learning/differentiable_ik/view_learned_yaw.py
+  - examples/learning/differentiable_fk/learn_pose_reaching.py
+  - examples/learning/differentiable_fk/view_pose_reaching.py
   watch_paths: [embodichain/compute/kinematics/, embodichain/learning/layers/,
     embodichain/learning/models/, embodichain/learning/__init__.py,
-    tests/compute/kinematics/, tests/learning/layers/, examples/learning/differentiable_ik/]
-  related_topics: [ik-solvers, rl-learning]
+    tests/compute/kinematics/, tests/learning/layers/, examples/learning/differentiable_ik/,
+    examples/learning/differentiable_fk/]
+  related_topics: [ik-solvers, rl-learning, supervised-learning]
   status: active
 - id: rl-learning
   title: RL Learning Pipeline

agent_context/topics/learning-layers/learning-layers.md

@@ -1,14 +1,20 @@
 # Differentiable learning layers
 
+The [Differentiable Kinematics guide](../../../docs/source/overview/learning/differentiable_kinematics.md)
+owns user-facing FK/IK setup, training examples and gradient/state contracts.
+
 ## Owners and dependency direction
 
 `compute/kinematics/{dls,pose_error,serial_chain,types}.py` owns fixed-step
 DLS, SE(3) residuals, copied PK geometry, and shared configuration/results.
 These modules do not import learning or simulation. `SerialChainKinematics`
 returns TCP poses and geometric Jacobians from one FK pass, retaining dJ/dq.
 
-`learning/layers/` owns `KinematicsLayer` and `DifferentiableIKLayer` Torch
-module interfaces. `learning/models/` owns general MLPs, feature normalization
+`learning/layers/fk.py` owns `DifferentiableFKLayer`: joints to root-frame TCP
+poses through a pure FK path; `evaluate` also returns geometric Jacobians from
+one FK pass. `DifferentiableIKLayer` composes this FK layer with fixed-step DLS.
+IK keeps its `kinematics` child name for checkpoint compatibility.
+`learning/models/` owns general MLPs, feature normalization
 and configuration-driven MLP construction. These shared components do not import
 RL, TensorDict or simulation. Concrete NeRD models, codecs and the shared-SL
 recipe live in the bundled `embodichain_models.dynamics.nerd` extension;
@@ -20,8 +26,8 @@ policy registration/builders belong to `learning/rl/policies/`; RL code consumes
 the shared networks and normalizer. See [RL learning](../rl-learning/rl-learning.md)
 for those contracts.
 
-Training, task-parameter optimization and the differentiable benchmark call
-`DifferentiableIKLayer` directly. Use `torch.no_grad()` for forward-only inference.
+Training, task-parameter optimization and the differentiable benchmark compose
+FK or IK layers directly. Use `torch.no_grad()` for forward-only inference.
 There is no registered Robot solver adapter: callers provide root-frame targets
 and explicitly configure the layer's TCP and limits. Existing Robot solver binding
 and arena conversion remain in [IK solvers](../ik-solvers/ik-solvers.md).
@@ -33,7 +39,9 @@ Supervised optimizer and offline-run ownership belongs to
 
 Targets are root-frame TCP matrices; seed batch axes match the target or contain
 only a joint axis. Training inputs must match model device and dtype. Both float32
-and float64 are supported. FK and Jacobian derivatives remain connected through
+and float64 are supported. FK inputs follow the model's joint order; FK does not
+project limits. Caller-owned pose losses can train joint predictors through FK.
+FK and Jacobian derivatives remain connected through
 all iterations. Final tolerances produce boolean metadata; no convergence-based
 Python exit occurs. Failure returns the final iterate and continuous residuals.
 Hard limits have piecewise gradients; disabling projection does not disable the
@@ -55,6 +63,13 @@ Model and layer callers own task losses and optimization. Formal task deployment
 belong to the task extension. See [RL learning](../rl-learning/rl-learning.md)
 for trainer/collector lifecycle.
 
+`examples/learning/differentiable_fk/learn_pose_reaching.py` trains an example-local
+joint predictor through fixed FK with TCP pose supervision and the shared SL
+trainer. Generating joint solutions are discarded; validation selects weights
+before independent local test evaluation. Its viewer compares initial and frozen
+trained predictions using one CPU worker and rejects stale target edits. Example
+smoke commands own launch/artifact checks; production tests own FK numerics.
+
 `examples/learning/differentiable_ik/learn_grasp_yaw.py` trains an example-local
 conditional yaw predictor through fixed-step IK, preserving grasp position and
 downward approach. Validation alone selects checkpoints; independent target

docs/source/api_reference/embodichain/embodichain.compute.rst

@@ -142,12 +142,14 @@ Geometry tensor kernels
 
    embodichain.compute.scene_geometry
 
-Differentiable inverse kinematics
---------------------------------
+Differentiable kinematics
+-------------------------
 
 The fixed-step DLS implementation preserves target and seed gradients through
 both FK and the geometric Jacobian. ``SerialChainKinematics`` adapts a copied
-PK chain without loading simulation. ``DifferentiableIKCfg`` configures the
+PK chain without loading simulation. Its ``forward_kinematics`` method computes
+only TCP poses, while ``evaluate`` also constructs differentiable Jacobians in
+the same FK pass. ``DifferentiableIKCfg`` configures the
 compute function and Torch layer; failed solves retain their last iterate.
 
 .. currentmodule:: embodichain.compute.kinematics

docs/source/api_reference/embodichain/embodichain.learning.rst

@@ -32,21 +32,24 @@ Differentiable layers
 ---------------------
 
 ``learning.layers`` composes simulation-independent computation as Torch modules.
-``KinematicsLayer`` owns fixed TCP and limit buffers; its copied PK geometry follows
+``DifferentiableFKLayer.forward`` maps joints to root-frame TCP transforms without
+constructing a geometric Jacobian; ``evaluate`` returns both in one FK pass.
+The layer owns fixed TCP and limit buffers, and 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.
+its geometry fingerprint is checked. ``DifferentiableIKLayer`` composes an FK layer:
+its ``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.
 
 .. currentmodule:: embodichain.learning.layers
 
 .. autosummary::
 
-   KinematicsLayer
+   DifferentiableFKLayer
    DifferentiableIKLayer
 
-.. automodule:: embodichain.learning.layers.kinematics
+.. automodule:: embodichain.learning.layers.fk
    :members:
 
 .. automodule:: embodichain.learning.layers.ik