From 1e160280ecca7bc6bf05750e4f833b75f45f80db Mon Sep 17 00:00:00 2001 From: yuecideng Date: Wed, 28 Oct 2026 16:59:27 +0000 Subject: [PATCH 1/3] feat(learning): add differentiable FK layer and reaching demo --- agent_context/MAP.yaml | 13 +- .../topics/learning-layers/learning-layers.md | 30 +- .../embodichain/embodichain.compute.rst | 8 +- .../embodichain/embodichain.learning.rst | 16 +- .../overview/learning/differentiable_ik.md | 118 +--- .../learning/differentiable_kinematics.md | 192 +++++++ docs/source/overview/learning/index.rst | 12 +- .../overview/sim/motion/solvers/index.rst | 2 +- .../compute/kinematics/serial_chain.py | 53 +- embodichain/learning/layers/__init__.py | 5 +- embodichain/learning/layers/fk.py | 180 +++++++ embodichain/learning/layers/ik.py | 6 +- embodichain/learning/layers/kinematics.py | 134 +---- examples/README.md | 5 + examples/learning/differentiable_fk/README.md | 107 ++++ .../differentiable_fk/learn_pose_reaching.py | 507 ++++++++++++++++++ .../differentiable_fk/requirements.txt | 5 + .../differentiable_fk/view_pose_reaching.py | 331 ++++++++++++ examples/learning/differentiable_ik/README.md | 6 + .../differentiable_ik/learn_grasp_yaw.py | 24 +- .../kinematic_solver/run_benchmark.py | 180 ++++++- tests/compute/kinematics/test_dls.py | 85 +++ tests/learning/layers/test_fk.py | 96 ++++ tests/learning/layers/test_ik.py | 18 +- tests/test_agent_context_map.py | 2 + 25 files changed, 1829 insertions(+), 306 deletions(-) create mode 100644 docs/source/overview/learning/differentiable_kinematics.md create mode 100644 embodichain/learning/layers/fk.py create mode 100644 examples/learning/differentiable_fk/README.md create mode 100644 examples/learning/differentiable_fk/learn_pose_reaching.py create mode 100644 examples/learning/differentiable_fk/requirements.txt create mode 100644 examples/learning/differentiable_fk/view_pose_reaching.py create mode 100644 tests/learning/layers/test_fk.py diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index b81e03e7f..c0705d8a0 100644 --- a/agent_context/MAP.yaml +++ b/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, KinematicsLayer, 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 diff --git a/agent_context/topics/learning-layers/learning-layers.md b/agent_context/topics/learning-layers/learning-layers.md index e1d5ff3c2..710eeaf0b 100644 --- a/agent_context/topics/learning-layers/learning-layers.md +++ b/agent_context/topics/learning-layers/learning-layers.md @@ -1,5 +1,8 @@ # 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 @@ -7,8 +10,12 @@ 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. +`KinematicsLayer` remains a direct alias at both its original module and package +import paths. 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 +27,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 +40,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 +64,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 @@ -69,4 +85,6 @@ Tests live under `tests/compute/kinematics/` and `tests/learning/layers/`. Validate FK, dJ/dq, full IK and recurrent-state gradients separately; include nonzero TCP offsets, repeated solves, checkpoint migration and import boundaries. GPU validation runs serially with `--run-gpu`. The unified kinematic benchmark's `differentiable` mode measures training and -forward-only baselines, including FK/Jacobian and DLS component costs. +forward-only baselines, including FK/Jacobian and DLS component costs. Its explicit +`fk` mode compares pure FK with pose extraction through `evaluate`, including +matching pose-loss gradients and separately measured warmed memory. diff --git a/docs/source/api_reference/embodichain/embodichain.compute.rst b/docs/source/api_reference/embodichain/embodichain.compute.rst index f7cfc46cc..706214743 100644 --- a/docs/source/api_reference/embodichain/embodichain.compute.rst +++ b/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 diff --git a/docs/source/api_reference/embodichain/embodichain.learning.rst b/docs/source/api_reference/embodichain/embodichain.learning.rst index 662af8072..17f6790dd 100644 --- a/docs/source/api_reference/embodichain/embodichain.learning.rst +++ b/docs/source/api_reference/embodichain/embodichain.learning.rst @@ -32,10 +32,14 @@ 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. ``KinematicsLayer`` remains a direct compatibility +alias of ``DifferentiableFKLayer``, including its original module import path. Both retain input gradients in training and evaluation mode. Callers own rollout state, optimizers, and detach boundaries. @@ -43,8 +47,12 @@ state, optimizers, and detach boundaries. .. autosummary:: - KinematicsLayer + DifferentiableFKLayer DifferentiableIKLayer + KinematicsLayer + +.. automodule:: embodichain.learning.layers.fk + :members: .. automodule:: embodichain.learning.layers.kinematics :members: diff --git a/docs/source/overview/learning/differentiable_ik.md b/docs/source/overview/learning/differentiable_ik.md index 9f5cb9ec6..346cfd2e3 100644 --- a/docs/source/overview/learning/differentiable_ik.md +++ b/docs/source/overview/learning/differentiable_ik.md @@ -1,113 +1,9 @@ -# Differentiable IK +--- +orphan: true +--- -Use `DifferentiableIKLayer` for training, task-parameter optimization and -forward-only inference. It calls the fixed-step DLS implementation in -`embodichain.compute.kinematics` without importing the simulation runtime. +# Differentiable FK and IK -## Training without simulation - -```python -from pathlib import Path -import torch -from embodichain.compute.kinematics import SerialChainKinematics, DifferentiableIKCfg -from embodichain.data import get_data_path -from embodichain.learning.layers import KinematicsLayer, DifferentiableIKLayer - -model = SerialChainKinematics.from_urdf( - Path(get_data_path("Franka/Panda/PandaWithHand.urdf")).read_text(), - end_link_name="fr3_hand_tcp", root_link_name="base", device="cuda", -) -kinematics = KinematicsLayer(model) -ik = DifferentiableIKLayer(kinematics, DifferentiableIKCfg(max_iterations=16)) -seed = kinematics.default_seed()[None].clone().requires_grad_() -target = kinematics(seed.detach() + 0.01).detach().requires_grad_() -result = ik.solve(target, seed) -loss = result.position_error.square().mean() + result.rotation_error.square().mean() -loss.backward() # Both seed.grad and target.grad are connected. -# q_des = ik(target, seed) returns just the joint tensor. -``` - -Targets are valid finite SE(3) transforms in the chain root frame and refer to the -configured TCP. Their shape is `(..., 4, 4)`. Seeds are `(dof,)` or `(..., dof)`; -results retain all leading dimensions. Positions use metres, revolute joints and -rotation vectors use radians, and prismatic coordinates use metres. - -The layer requires matching dtype/device and supports float32 and float64. A model -copy owns the geometry; `layer.to(...)` migrates its PK tensors as well as registered -buffers. Checkpoints contain numerical settings, TCP, limits and a geometry -fingerprint. Reconstruct the same chain before loading. Geometry and TCP are fixed -configuration, not learned parameters. `eval()` does not disable gradients. - -## Forward-only inference and frame ownership - -```python -with torch.no_grad(): - result = ik.solve(target, seed) - success, joints = result.success, result.qpos -``` - -The layer takes one seed per target and does not run a multi-start search. It is -an `nn.Module`, not a registered Robot solver. Callers own world-to-root frame -conversion and must configure TCP and physical joint limits explicitly when -constructing `KinematicsLayer`. Existing simulation solvers and robot defaults -remain unchanged. - -## Numerical and gradient contract - -Each iteration computes FK and the TCP Jacobian once, then applies -`q_next = q + step_size * J.T @ solve(J @ J.T + damping**2 * I, error)`. -FK uses standard autograd, and the derivative of the Jacobian with respect to joints -is retained. The implementation uses `solve_ex(check_errors=False)` without -per-iteration convergence synchronization. No previous seed is stored across calls. - -Defaults are 16 iterations, damping 0.05, step size 0.5, position tolerance 0.0005 m -and rotation tolerance 0.0005 rad. Tolerances classify the final iterate, never stop -iterations early. Hard limit projection applies to the seed and every update; -outside a clamp interval its derivative is zero. Set `enforce_joint_limits=False` -for unconstrained experiments, and supply your own limit loss; success still checks -physical limits. Unbounded joints default to projected zero; bounded joints use -their midpoint. - -Rotation errors use shortest root-frame axis-angle vectors with a finite derivative -at zero. The pi branch cut and hard limit transitions are not globally smooth. -Only first-order input gradients are supported as a public contract. DLS is a -local solver, not a guarantee of global reachability. Failed solves return the last -iterate and differentiable residuals rather than substituting the seed. - -## Measurement - -```bash -python -m scripts.benchmark.robotics.kinematic_solver.run_benchmark -s differentiable -``` - -The benchmark measures forward, forward+backward, FK/Jacobian and DLS components on -identical local and distant Franka targets. FEP and PytorchSolver are forward-only -baselines. It reports errors including failed rows, throughput, CPU memory and -Torch GPU allocations; Warp allocations are not included in that GPU counter. - -## Learning grasp orientation - -The standalone example trains a conditional yaw predictor through the IK layer. -It receives current joints and grasp position, predicts a yaw correction, and -uses joint displacement plus pose residuals to train shared network parameters. -Frozen checkpoints can predict new goals without per-goal gradient optimization. -The viewer compares heading, a learned predictor and a numerical reference. - -```bash -conda activate py311 -python -m examples.learning.differentiable_ik.learn_grasp_yaw --device cuda -python -m examples.learning.differentiable_ik.view_learned_yaw --port 8101 -``` - -The experiment compares a shared constant, IK-trained and teacher-supervised -predictors, online optimization and an angular grid. Independent validation -selects checkpoints; IID, disjoint spatial and initial-state-shift sets measure -transfer. Training uses 24-step IK and every final result is checked with 64 -steps. Reports include teacher-label cost and complete inference latency. - -See `examples/learning/differentiable_ik/README.md` for training and preview commands, -measured results, gradient checks, dependencies and exact evaluation protocol. -The first training seed supplies the viewer; frozen inference runs on one CPU -worker and stale results are discarded after input changes. Numerical angle optimization and grid search are integrated reference baselines -in the same training/evaluation entry point. -Collision, contact, dynamics and physical execution are outside this example. +This guide has moved to {doc}`differentiable_kinematics`, which covers +shared setup, FK and IK training, gradient/state contracts, examples and +measurement. diff --git a/docs/source/overview/learning/differentiable_kinematics.md b/docs/source/overview/learning/differentiable_kinematics.md new file mode 100644 index 000000000..6a0f672d9 --- /dev/null +++ b/docs/source/overview/learning/differentiable_kinematics.md @@ -0,0 +1,192 @@ +# Differentiable Kinematics + +Use robot kinematics inside a Torch network or task loss without starting a +simulation environment. `DifferentiableFKLayer` maps joints to TCP poses; +`DifferentiableIKLayer` composes the FK layer with fixed-step damped least squares +(DLS) to map target TCP poses to joints. Both support supervised learning, +task-parameter optimization and forward-only inference. + +| Call | Input | Output | Input gradients | +| --- | --- | --- | --- | +| `fk(qpos)` / `fk.forward(qpos)` | Joints `(..., dof)` | TCP poses `(..., 4, 4)` | Joint positions | +| `fk.evaluate(qpos)` | Joints `(..., dof)` | TCP poses and geometric Jacobians `(..., 6, dof)` | Joint positions, including Jacobian derivatives | +| `fk.jacobian(qpos)` | Joints `(..., dof)` | Geometric Jacobians `(..., 6, dof)` | Joint positions | +| `ik(target_xpos, qpos_seed)` / `ik.forward(...)` | TCP targets and optional joint seeds | Joints `(..., dof)` | Targets and supplied seeds | +| `ik.solve(target_xpos, qpos_seed)` | Same as `ik.forward` | Final joints, success mask and pose residuals | Floating outputs retain target and seed gradients; the mask is boolean | + +The layers are standalone `nn.Module` components. For simulator-backed Robot +solving and differential pose control, see the +{doc}`motion solvers `. Full signatures and +configuration fields are in the +{doc}`learning API ` and +{doc}`kinematics compute API `. + +## Shared initialization + +The examples below use the same fixed Franka model on CPU. Run this initialization +before either the FK or IK example; the two examples do not depend on each other. + +```python +from pathlib import Path + +import torch + +from embodichain.compute.kinematics import ( + DifferentiableIKCfg, + SerialChainKinematics, + pose_error, +) +from embodichain.data import get_data_path +from embodichain.learning.layers import DifferentiableFKLayer, DifferentiableIKLayer + +model = SerialChainKinematics.from_urdf( + Path(get_data_path("Franka/Panda/PandaWithHand.urdf")).read_text(), + root_link_name="base", + end_link_name="fr3_hand_tcp", + device="cpu", + dtype=torch.float32, +) +tcp = torch.eye(4, device="cpu", dtype=torch.float32) +fk = DifferentiableFKLayer(model, tcp=tcp) +``` + +`end_link_name` selects the model's terminal frame. `tcp` is the fixed transform +from that frame to the task's tool frame; identity uses the terminal frame +directly. The FK layer copies the model and owns its TCP and physical-limit +buffers. Configure a different TCP or physical limits at construction when the +task requires them. + +## Differentiable FK + +Train a joint predictor using TCP pose supervision. This example creates a small +predictor and illustrative reachable labels, then backpropagates position and +rotation errors through FK into the predictor's parameters. Replace the example +observations and labels with task data for training. + +```python +predictor = torch.nn.Linear( + 3, fk.dof, device=fk.tcp.device, dtype=fk.tcp.dtype +) +observation = fk.tcp.new_ones(8, 3) +with torch.no_grad(): + target_pose = fk(fk.default_seed().expand(8, -1)) + +q_pred = fk.default_seed() + 0.1 * predictor(observation) +pose_pred = fk(q_pred) +position_error, rotation_error = pose_error(pose_pred, target_pose) +loss = ( + position_error.square().sum(-1) + + rotation_error.square().sum(-1) +).mean() +loss.backward() # Gradients reach predictor.weight and predictor.bias. +``` + +`fk(qpos)` computes poses without constructing an explicit geometric Jacobian. +Use `fk.evaluate(qpos)` when the loss also needs Jacobians: it returns both from +one FK pass. Geometric Jacobians are expressed in the chain root frame, with +linear rows preceding angular rows. + +FK evaluates the supplied joints without clipping them to physical limits. The +training task owns loss weights, joint-limit handling and any posture loss. +TCP pose supervision alone need not select a unique joint posture for a redundant +robot; use an output mapping or additional objective when the task needs one. + +## Differentiable IK + +IK takes a target pose and an optional starting joint configuration. The example +below creates its own seed and nearby target, then differentiates a pose-residual +loss through the unrolled solver. + +```python +ik = DifferentiableIKLayer(fk, cfg=DifferentiableIKCfg(max_iterations=16)) +qpos_seed = fk.default_seed()[None].clone().requires_grad_() +target_pose = fk(qpos_seed.detach() + 0.01).detach().requires_grad_() + +result = ik.solve(target_pose, qpos_seed=qpos_seed) +loss = ( + result.position_error.square().mean() + + result.rotation_error.square().mean() +) +loss.backward() # Both qpos_seed.grad and target_pose.grad are connected. +``` + +Use `ik(target_pose, qpos_seed=...)` when a network needs only the joint tensor. +`ik.solve(...)` additionally returns a boolean `success` mask and continuous +position/rotation residuals. Seeds are `(dof,)` or match all target batch axes; +omitting a seed uses bounded joint midpoints, or projected zero for unbounded +joints. There is one seed per target and no multi-start search. + +Each DLS iteration computes FK and the TCP Jacobian once, retaining their +derivatives, then applies: + +```text +q_next = q + step_size * J.T @ solve(J @ J.T + damping**2 * I, error) +``` + +The iteration count is fixed. Position and rotation tolerances classify the final +iterate without stopping the loop early. `success` also requires finite values +and physical-limit validity. A failed solve returns the last iterate and its +differentiable residuals. DLS is a local solver; it does not guarantee global +reachability. + +Joint-limit projection applies to the seed and every update by default. Its +gradient is zero outside a clamp interval. Set `enforce_joint_limits=False` for +an unconstrained objective with caller-owned limit handling; the final success +check still uses physical limits. Numerical options belong to +`DifferentiableIKCfg` in the compute API. + +## Shared gradient and state contracts + +- **Frames and units:** poses are valid finite SE(3) matrices in the configured + chain root frame and refer to the configured TCP. Callers own world-to-root + conversion. Joint inputs follow `model.joint_names`. Translation and prismatic + coordinates use metres; revolute joints and rotation vectors use radians. +- **Batch, device and precision:** FK joints are `(..., dof)` and poses are + `(..., 4, 4)`; leading batch dimensions are retained. Inputs must match the + layer's device and dtype. Float32 and float64 are supported. `.to(...)` migrates + copied PK geometry as well as registered buffers. +- **Gradients:** only first-order input gradients are part of the public + contract. Geometry and TCP are fixed configuration, not trained parameters. + `pose_error` returns translation and shortest root-frame rotation vectors; + the zero-rotation derivative is finite, while the shortest rotation has a + branch cut at pi. `.eval()` preserves gradients; use `torch.no_grad()` for + forward-only inference. +- **Caller-owned state:** neither layer owns an optimizer or recurrent seed. + Training code owns losses, rollout state, reset and detach boundaries. + Outstanding calls retain independent computation graphs. +- **Checkpoints:** reconstruct the same robot geometry before loading a state + dict. Checkpoints store TCP/limit buffers, geometry identity and, for IK, + numerical settings. `load_state_dict(assign=True)` also migrates fixed + geometry to the loaded buffers' device and dtype, including nested loads. + `KinematicsLayer` remains a direct alias of `DifferentiableFKLayer` at its + original package and module import paths. IK retains the `kinematics` + constructor argument and child name for existing state dicts. + +## Examples and measurement + +For complete training and browser previews, see the example READMEs: + +- [FK pose reaching](https://github.com/DexForce/EmbodiChain/blob/main/examples/learning/differentiable_fk/README.md) + trains a joint predictor through TCP pose loss and compares initial and learned + predictions on held-out cases and edited targets. +- [IK grasp orientation](https://github.com/DexForce/EmbodiChain/blob/main/examples/learning/differentiable_ik/README.md) + trains a conditional yaw predictor through IK and compares frozen predictions + against heading and numerical references. + +The READMEs own runnable commands, dependencies, artifacts and evaluation +protocols. These are kinematic examples; collision, contact, dynamics and physical +execution are outside their validation scope. + +The unified kinematic benchmark provides two relevant modes: + +```bash +python -m scripts.benchmark.robotics.kinematic_solver.run_benchmark -s fk +python -m scripts.benchmark.robotics.kinematic_solver.run_benchmark -s differentiable +``` + +`fk` compares pure FK with pose extraction through `evaluate` using identical +inputs and pose losses, checking pose/gradient agreement and separately measuring +warmed memory. `differentiable` measures IK forward and forward/backward passes, +FK/Jacobian and DLS components, and FEP/Pytorch forward-only baselines. Reports +describe workloads, error metrics, timing and memory-accounting scope; benchmark +results do not establish performance for a different task or device. diff --git a/docs/source/overview/learning/index.rst b/docs/source/overview/learning/index.rst index 4525f3b6b..88fdc82b1 100644 --- a/docs/source/overview/learning/index.rst +++ b/docs/source/overview/learning/index.rst @@ -2,16 +2,18 @@ Learning Infrastructure ======================= Embodied learning trains agents to act through a robot's observations and -actions. This section covers policy training and evaluation in EmbodiChain; -the built-in training framework currently focuses on reinforcement learning. -Standalone differentiable kinematics layers also support task-parameter -optimization and training without a simulation environment. +actions. This section describes the learning infrastructure and its policy +training and evaluation workflows. +Shared Torch models live in ``learning.models``; ``learning.layers`` composes +simulation-independent kinematics from ``compute.kinematics`` as reusable network +layers. For FK/IK usage, pose losses and gradient contracts, see +:doc:`differentiable_kinematics`. .. toctree:: :maxdepth: 1 /overview/rl/index - differentiable_ik + differentiable_kinematics Related Workflows ----------------- diff --git a/docs/source/overview/sim/motion/solvers/index.rst b/docs/source/overview/sim/motion/solvers/index.rst index 9694191f6..e8296ae44 100644 --- a/docs/source/overview/sim/motion/solvers/index.rst +++ b/docs/source/overview/sim/motion/solvers/index.rst @@ -37,7 +37,7 @@ Choose a solver Analytical solvers require matching robot geometry. Numerical solvers support more general chains but depend on seed quality and convergence settings. For differentiable training and task-parameter optimization, use the standalone -:doc:`Torch IK layer `. +:doc:`differentiable FK / IK layers `. .. _motion-solver-conventions: diff --git a/embodichain/compute/kinematics/serial_chain.py b/embodichain/compute/kinematics/serial_chain.py index 6b75d6223..25a3359fe 100644 --- a/embodichain/compute/kinematics/serial_chain.py +++ b/embodichain/compute/kinematics/serial_chain.py @@ -108,19 +108,10 @@ def to( """Return a model copy on the requested device and precision.""" return type(self)(self.chain, device, dtype) - def evaluate( + def _prepare_inputs( self, qpos: torch.Tensor, tcp: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: - """Compute TCP transforms and geometric Jacobians using one FK pass. - - Args: - qpos: Joint positions, shape ``(..., dof)``. - tcp: Flange-to-tool transforms, shape ``(4, 4)`` or ``(..., 4, 4)``. - - Returns: - TCP poses ``(..., 4, 4)`` and root-frame Jacobians ``(..., 6, dof)``. - Their backward paths include both FK and Jacobian derivatives. - """ + """Validate shared FK/Jacobian inputs and flatten leading batch axes.""" if ( qpos.ndim < 1 or qpos.shape[-1] != self.dof @@ -134,13 +125,49 @@ def evaluate( or tcp.dtype != qpos.dtype ): raise ValueError("qpos, TCP and kinematics model must share dtype/device") + return qpos.reshape(-1, self.dof), tcp.reshape(-1, 4, 4) + + def forward_kinematics(self, qpos: torch.Tensor, tcp: torch.Tensor) -> torch.Tensor: + """Compute TCP transforms without constructing geometric Jacobians. + + Args: + qpos: Joint positions in ``joint_names`` order, shape ``(..., dof)``. + tcp: Flange-to-tool transforms, shape ``(4, 4)`` or ``(..., 4, 4)``. + Both tensors must share this model's device and dtype. + + Returns: + Root-frame TCP poses, shape ``(..., 4, 4)``, retaining input + gradients. Joint limits are not enforced by FK. + + Raises: + ValueError: If input shapes, devices or dtypes do not match. + """ + q, tool_offset = self._prepare_inputs(qpos, tcp) + transforms = self.chain.forward_kinematics_tensor(q, analytical_grad=False) + flange = transforms[self.chain._serial_eef_frame_idx] + tool = flange @ tool_offset + return tool.reshape(*qpos.shape[:-1], 4, 4) + + def evaluate( + self, qpos: torch.Tensor, tcp: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """Compute TCP transforms and geometric Jacobians using one FK pass. + + Args: + qpos: Joint positions, shape ``(..., dof)``. + tcp: Flange-to-tool transforms, shape ``(4, 4)`` or ``(..., 4, 4)``. + + Returns: + TCP poses ``(..., 4, 4)`` and root-frame Jacobians ``(..., 6, dof)``. + Their backward paths include both FK and Jacobian derivatives. + """ + q, tool_offset = self._prepare_inputs(qpos, tcp) batch = qpos.shape[:-1] - q = qpos.reshape(-1, self.dof) transforms = self.chain.forward_kinematics_tensor(q, analytical_grad=False) jacobian, flange = self.chain.jacobian_tensor( q, ret_eef_pose=True, all_transforms=transforms ) - tool = flange @ tcp.reshape(-1, 4, 4) + tool = flange @ tool_offset offset = tool[:, :3, 3] - flange[:, :3, 3] correction = torch.cross( jacobian[:, 3:, :].transpose(-1, -2), diff --git a/embodichain/learning/layers/__init__.py b/embodichain/learning/layers/__init__.py index cef007d71..96530a36c 100644 --- a/embodichain/learning/layers/__init__.py +++ b/embodichain/learning/layers/__init__.py @@ -18,7 +18,8 @@ from __future__ import annotations -from .kinematics import KinematicsLayer +from .fk import DifferentiableFKLayer from .ik import DifferentiableIKLayer +from .kinematics import KinematicsLayer -__all__ = ["KinematicsLayer", "DifferentiableIKLayer"] +__all__ = ["DifferentiableFKLayer", "DifferentiableIKLayer", "KinematicsLayer"] diff --git a/embodichain/learning/layers/fk.py b/embodichain/learning/layers/fk.py new file mode 100644 index 000000000..ca07cd269 --- /dev/null +++ b/embodichain/learning/layers/fk.py @@ -0,0 +1,180 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""A differentiable forward-kinematics layer with fixed robot geometry.""" + +from __future__ import annotations + +from typing import Any + +import torch +from torch import nn + +from embodichain.compute.kinematics.serial_chain import SerialChainKinematics + +__all__ = ["DifferentiableFKLayer"] + + +class DifferentiableFKLayer(nn.Module): + """Expose FK and differentiable Jacobians as a reusable network layer. + + Args: + model: Fixed kinematics model, copied so device moves are independent. + tcp: Optional flange-to-tool rigid transform. + lower: Optional physical lower limits; defaults to the chain limits. + upper: 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. + Joint inputs follow ``model.joint_names`` and are not clipped to limits. + Only first-order input gradients are part of the public contract. + """ + + def __init__( + self, + model: SerialChainKinematics, + tcp: torch.Tensor | None = None, + lower: torch.Tensor | None = None, + upper: torch.Tensor | None = None, + ) -> None: + super().__init__() + self.model = model.to(model.chain.device, model.chain.dtype) + options = {"device": model.chain.device, "dtype": model.chain.dtype} + tcp_tensor = ( + torch.eye(4, **options) + if tcp is None + else torch.as_tensor(tcp, **options).detach().clone() + ) + if tcp_tensor.shape != (4, 4) or not torch.isfinite(tcp_tensor).all(): + raise ValueError("TCP must be a finite 4x4 rigid transform") + rotation = tcp_tensor[:3, :3] + if ( + not torch.allclose( + rotation.T @ rotation, torch.eye(3, **options), atol=1e-5, rtol=1e-5 + ) + or not torch.allclose( + torch.linalg.det(rotation), torch.ones((), **options), atol=1e-5 + ) + or not torch.allclose(tcp_tensor[3], torch.tensor([0, 0, 0, 1], **options)) + ): + raise ValueError("TCP must be a rigid SE(3) transform") + lo = ( + torch.as_tensor(model.chain.low if lower is None else lower, **options) + .detach() + .clone() + ) + hi = ( + torch.as_tensor(model.chain.high if upper is None else upper, **options) + .detach() + .clone() + ) + if ( + lo.shape != (model.dof,) + or hi.shape != lo.shape + or torch.isnan(lo).any() + or torch.isnan(hi).any() + or (lo > hi).any() + ): + raise ValueError("Invalid physical joint limits") + self.register_buffer("tcp", tcp_tensor) + self.register_buffer("lower", lo) + self.register_buffer("upper", hi) + self.register_load_state_dict_post_hook(self._sync_model_after_load) + + @property + def dof(self) -> int: + """Number of movable joints in solver order.""" + return self.model.dof + + def _apply(self, fn: Any, recurse: bool = True) -> DifferentiableFKLayer: + # PK is not an nn.Module: explicitly migrate all of its geometry tensors. + probe = fn(self.tcp) + if probe.dtype not in (torch.float32, torch.float64): + raise ValueError("DifferentiableFKLayer supports float32 and float64") + super()._apply(fn, recurse=recurse) + self.model = self.model.to(self.tcp.device, self.tcp.dtype) + return self + + def _sync_model_after_load( + self, _module: nn.Module, _incompatible_keys: Any + ) -> None: + # assign=True replaces buffers without calling _apply, also when a parent + # network loads this layer recursively. PK geometry is not a buffer. + if self.tcp.dtype not in (torch.float32, torch.float64): + raise ValueError("DifferentiableFKLayer supports float32 and float64") + if ( + self.model.chain.device != self.tcp.device + or self.model.chain.dtype != self.tcp.dtype + ): + self.model = self.model.to(self.tcp.device, self.tcp.dtype) + + def get_extra_state(self) -> dict[str, str]: + """Record fixed geometry identity alongside tensor buffers.""" + return {"model_fingerprint": self.model.fingerprint} + + def set_extra_state(self, state: dict[str, str]) -> None: + """Reject checkpoints belonging to a different kinematic model.""" + if state != self.get_extra_state(): + raise ValueError( + "Kinematics checkpoint geometry does not match the configured model" + ) + + def forward(self, qpos: torch.Tensor) -> torch.Tensor: + """Compute TCP transforms without constructing a geometric Jacobian. + + Args: + qpos: Joint positions in model order, shape ``(..., dof)``. Device + and dtype must match this layer's buffers and geometry. + + Returns: + Root-frame TCP transforms, shape ``(..., 4, 4)``, retaining input + gradients and all leading batch dimensions. + """ + return self.model.forward_kinematics(qpos, self.tcp) + + def evaluate(self, qpos: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Compute TCP transforms and geometric Jacobians in one FK pass. + + Args: + qpos: Joint positions in model order, shape ``(..., dof)``. + + Returns: + Root-frame TCP transforms ``(..., 4, 4)`` and geometric Jacobians + ``(..., 6, dof)``, with linear rows preceding angular rows. + """ + return self.model.evaluate(qpos, self.tcp) + + def jacobian(self, qpos: torch.Tensor) -> torch.Tensor: + """Return differentiable TCP Jacobians in the root frame. + + Args: + qpos: Joint positions in model order, shape ``(..., dof)``. + + Returns: + Geometric Jacobians ``(..., 6, dof)``, with linear rows preceding + angular rows. Their gradients remain connected to ``qpos``. + """ + return self.evaluate(qpos)[1] + + def default_seed(self) -> torch.Tensor: + """Return bounded joint midpoints, or projected zero for unbounded joints.""" + bounded = torch.isfinite(self.lower) & torch.isfinite(self.upper) + midpoint = (self.lower.where(bounded, 0) + self.upper.where(bounded, 0)) * 0.5 + return midpoint.clamp(min=self.lower, max=self.upper) diff --git a/embodichain/learning/layers/ik.py b/embodichain/learning/layers/ik.py index cb763e5ec..86d742911 100644 --- a/embodichain/learning/layers/ik.py +++ b/embodichain/learning/layers/ik.py @@ -30,7 +30,7 @@ DifferentiableIKResult, solve_dls_ik, ) -from .kinematics import KinematicsLayer +from .fk import DifferentiableFKLayer __all__ = ["DifferentiableIKLayer"] @@ -39,7 +39,7 @@ class DifferentiableIKLayer(nn.Module): """Compose fixed-step DLS with a differentiable kinematics module. Args: - kinematics: Model owning geometry, TCP and physical limit buffers. + kinematics: FK layer owning geometry, TCP and physical limit buffers. cfg: Numerical settings, copied at construction. Note: @@ -48,7 +48,7 @@ class DifferentiableIKLayer(nn.Module): """ def __init__( - self, kinematics: KinematicsLayer, cfg: DifferentiableIKCfg | None = None + self, kinematics: DifferentiableFKLayer, cfg: DifferentiableIKCfg | None = None ) -> None: super().__init__() self.kinematics = kinematics diff --git a/embodichain/learning/layers/kinematics.py b/embodichain/learning/layers/kinematics.py index 4701f5eea..83dab9cbe 100644 --- a/embodichain/learning/layers/kinematics.py +++ b/embodichain/learning/layers/kinematics.py @@ -14,140 +14,12 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Torch module interface to fixed robot geometry, independent of simulation.""" +"""Compatibility name for :class:`.fk.DifferentiableFKLayer`.""" from __future__ import annotations -from typing import Any - -import torch -from torch import nn - -from embodichain.compute.kinematics.serial_chain import SerialChainKinematics +from .fk import DifferentiableFKLayer __all__ = ["KinematicsLayer"] - -class KinematicsLayer(nn.Module): - """Expose FK and differentiable Jacobians as a reusable network layer. - - Args: - model: Fixed kinematics model, copied so device moves are independent. - tcp: Optional flange-to-tool rigid transform. - lower: Optional physical lower limits; defaults to the chain limits. - upper: 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. - """ - - def __init__( - self, - model: SerialChainKinematics, - tcp: torch.Tensor | None = None, - lower: torch.Tensor | None = None, - upper: torch.Tensor | None = None, - ) -> None: - super().__init__() - self.model = model.to(model.chain.device, model.chain.dtype) - options = {"device": model.chain.device, "dtype": model.chain.dtype} - tcp_tensor = ( - torch.eye(4, **options) - if tcp is None - else torch.as_tensor(tcp, **options).detach().clone() - ) - if tcp_tensor.shape != (4, 4) or not torch.isfinite(tcp_tensor).all(): - raise ValueError("TCP must be a finite 4x4 rigid transform") - rotation = tcp_tensor[:3, :3] - if ( - not torch.allclose( - rotation.T @ rotation, torch.eye(3, **options), atol=1e-5, rtol=1e-5 - ) - or not torch.allclose( - torch.linalg.det(rotation), torch.ones((), **options), atol=1e-5 - ) - or not torch.allclose(tcp_tensor[3], torch.tensor([0, 0, 0, 1], **options)) - ): - raise ValueError("TCP must be a rigid SE(3) transform") - lo = ( - torch.as_tensor(model.chain.low if lower is None else lower, **options) - .detach() - .clone() - ) - hi = ( - torch.as_tensor(model.chain.high if upper is None else upper, **options) - .detach() - .clone() - ) - if ( - lo.shape != (model.dof,) - or hi.shape != lo.shape - or torch.isnan(lo).any() - or torch.isnan(hi).any() - or (lo > hi).any() - ): - raise ValueError("Invalid physical joint limits") - self.register_buffer("tcp", tcp_tensor) - self.register_buffer("lower", lo) - self.register_buffer("upper", hi) - self.register_load_state_dict_post_hook(self._sync_model_after_load) - - @property - def dof(self) -> int: - """Number of movable joints in solver order.""" - return self.model.dof - - def _apply(self, fn: Any, recurse: bool = True) -> KinematicsLayer: - # PK is not an nn.Module: explicitly migrate all of its geometry tensors. - probe = fn(self.tcp) - if probe.dtype not in (torch.float32, torch.float64): - raise ValueError("KinematicsLayer supports float32 and float64") - super()._apply(fn, recurse=recurse) - self.model = self.model.to(self.tcp.device, self.tcp.dtype) - return self - - def _sync_model_after_load( - self, _module: nn.Module, _incompatible_keys: Any - ) -> None: - # assign=True replaces buffers without calling _apply, also when a parent - # network loads this layer recursively. PK geometry is not a buffer. - if self.tcp.dtype not in (torch.float32, torch.float64): - raise ValueError("KinematicsLayer supports float32 and float64") - if ( - self.model.chain.device != self.tcp.device - or self.model.chain.dtype != self.tcp.dtype - ): - self.model = self.model.to(self.tcp.device, self.tcp.dtype) - - def get_extra_state(self) -> dict[str, str]: - """Record fixed geometry identity alongside tensor buffers.""" - return {"model_fingerprint": self.model.fingerprint} - - def set_extra_state(self, state: dict[str, str]) -> None: - """Reject checkpoints belonging to a different kinematic model.""" - if state != self.get_extra_state(): - raise ValueError( - "Kinematics checkpoint geometry does not match the configured model" - ) - - def forward(self, qpos: torch.Tensor) -> torch.Tensor: - """Return root-frame TCP transforms for ``(..., dof)`` joints.""" - return self.evaluate(qpos)[0] - - def evaluate(self, qpos: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """Return TCP transforms and root-frame geometric Jacobians.""" - return self.model.evaluate(qpos, self.tcp) - - def jacobian(self, qpos: torch.Tensor) -> torch.Tensor: - """Return differentiable TCP Jacobians, shape ``(..., 6, dof)``.""" - return self.evaluate(qpos)[1] - - def default_seed(self) -> torch.Tensor: - """Return bounded joint midpoints, or projected zero for unbounded joints.""" - bounded = torch.isfinite(self.lower) & torch.isfinite(self.upper) - midpoint = (self.lower.where(bounded, 0) + self.upper.where(bounded, 0)) * 0.5 - return midpoint.clamp(min=self.lower, max=self.upper) +KinematicsLayer = DifferentiableFKLayer diff --git a/examples/README.md b/examples/README.md index 31b1e903b..d6840b972 100644 --- a/examples/README.md +++ b/examples/README.md @@ -44,6 +44,11 @@ The [SL data-source example](learning/sl_data_source/README.md) connects residen causal windows to explicit action or state-transition encoders, and describes caller-owned online sampling and the limits of exact recovery. +The [differentiable FK example](learning/differentiable_fk/README.md) trains a local +Franka joint predictor through fixed FK using TCP pose supervision and the shared +SL trainer. Its browser preview compares initialized and trained predictions on +held-out targets, with frozen inference for interactive target edits. + The [RWM fitting example](learning/rwm/README.md) qualifies controlled supervised state-transition fitting and state/action gradients. Its generated recipes use the installable `embodichain_models.dynamics.rwm` model family and shared SL diff --git a/examples/learning/differentiable_fk/README.md b/examples/learning/differentiable_fk/README.md new file mode 100644 index 000000000..35be6100d --- /dev/null +++ b/examples/learning/differentiable_fk/README.md @@ -0,0 +1,107 @@ +# Learn reaching through differentiable FK + +Train a small Franka joint predictor using **TCP pose labels**. A fixed +`DifferentiableFKLayer` maps its predicted joints to TCP transforms; position and +rotation errors backpropagate through FK into the network. The learned model is a +local approximate inverse mapping. Robot geometry stays fixed. + +```text +current joints + target TCP → MLP → predicted joints → FK → pose loss + ↑ │ + └──────── backward ───────────┘ +``` + +The example uses `learning.models.MLP` and the common `TensorBatchSource`, +`SupervisedLearner` and `SupervisedTrainer`. It needs no simulator, IK solver, +physics backend or target-joint labels. Models and task losses remain example-local. + +## Train and view + +From a checkout with EmbodiChain installed: + +```bash +pip install -r examples/learning/differentiable_fk/requirements.txt + +embodichain run-example learning/differentiable_fk/learn_pose_reaching \ + --device cpu --output outputs/learn_fk_reaching +embodichain run-example learning/differentiable_fk/view_pose_reaching \ + --run outputs/learn_fk_reaching --port 8103 +``` + +Open `http://127.0.0.1:8103`; forward the port when running on a remote host. +The blue Franka shows the zero-initialized network's output, equal to the current +posture. The green Franka shows one frozen prediction from the validation-selected +model. Each robot has its own copy of the same root-frame target. Large axes mark +the goal, small axes mark the achieved TCP, and yellow lines show position error. +The sidebar reports position/rotation errors, independent validation curves, and +before/after metrics over **all** held-out test samples. + +Choose a test case or edit its XYZ and roll/pitch/yaw offsets. Reset restores the +saved target; changing the case also clears offsets. Rotation edits compose +`Rz(yaw) @ Ry(pitch) @ Rx(roll)` in the chain root frame. Edited targets are labeled +separately from the saved evaluation, may be unreachable, and may exceed the +learned local domain. Their actual errors remain visible. The viewer uses one CPU +worker, coalesces requests and rejects stale results; it never updates weights. + +`python -m examples.learning.differentiable_fk.learn_pose_reaching` and +`python -m examples.learning.differentiable_fk.view_pose_reaching` are equivalent +checkout entry points. Training also accepts `--device cuda`, `--seed` and +`--updates`; the viewer accepts `--host`, `--port`, and `--run`. + +## Learning contract + +- The robot chain is `base → fr3_hand_tcp` from the official Franka URDF. + Positions use metres; joints and rotation residuals use radians. +- Sample current joints in the intersection of the official ready posture + ±0.10 rad and physical limits with a 0.15 rad margin. Generate target joints by + adding independent ±0.10 rad perturbations, apply FK, and discard those target + joints. Training and saved datasets contain only current joints and TCP poses. +- Input features have 13 dimensions: seven current-joint offsets from ready, + three current-to-target position errors and three root-frame rotation errors. + Their respective scales are 0.10 rad, 0.05 m and 0.20 rad. +- The MLP is `13 → 64 → 64 → 7` with SiLU. Its final layer starts at zero. + Predictions are `q_current + 0.15 * tanh(raw)`. The admitted current-state + margin keeps every prediction within physical limits. +- Loss is the mean squared norm of position/0.05 and rotation/0.20 residuals, + plus `1e-3 * mean(||q_pred - q_current||²)`. There is no unique ground-truth + joint solution for a redundant arm; the posture term favors nearby solutions. + +Defaults are float32, CPU, 1,024 training examples, 128 validation examples and +128 independent test examples. Model seed is 17; corpus seeds are 23/24/25 and +the SL sampling seed is model seed + 12. Adam uses learning rate 0.001, batch 128 +and 1,000 updates. Validation every 50 updates selects the checkpoint. Test data +never selects weights, feature scales or hyperparameters. These held-outs measure +local interpolation, not global workspace generalization or multiple-seed robustness. + +## Artifacts and checks + +The output contains common-SL `run.json`, `metrics.jsonl` and `last.pt`, plus: + +- `predictor.pt`: exported validation-selected weights, FK identity and config; +- `datasets.npz`: exact train/validation/test inputs and TCP labels; +- `evaluation.npz`: every test prediction, position/rotation error and success flag; +- `metrics.json`, `report.md`, `learning.png`: full-test summaries and validation curves. + +The default complete run requires held-out position and rotation mean errors to +each improve by at least 50%. Reports include means, P95 and the fraction meeting +**both 2 mm and 2°**, retaining every failed sample. Execution checks verify +pose-only parameter gradients, finite bounded outputs, unchanged FK state and +CPU-reloaded prediction agreement. Use a fresh output directory for a new run; +the common trainer protects existing run history. + +For a bounded execution check: + +```bash +embodichain run-example learning/differentiable_fk/learn_pose_reaching \ + --smoke --output outputs/learn_fk_reaching_smoke +embodichain run-example learning/differentiable_fk/view_pose_reaching \ + --run outputs/learn_fk_reaching_smoke --port 8104 --smoke +``` + +Training smoke uses 20 updates and 64/16/16 samples. It checks the pipeline without +requiring the learning-improvement gate. Viewer smoke constructs both robot +scenes, edits one target, checks bounded frozen predictions and HTTP, then exits. +Production FK tests own the numerical and general gradient contracts. + +The preview shows kinematic configurations. It does not execute a trajectory or +validate collisions, contacts, dynamics or hardware control. diff --git a/examples/learning/differentiable_fk/learn_pose_reaching.py b/examples/learning/differentiable_fk/learn_pose_reaching.py new file mode 100644 index 000000000..c0c091a08 --- /dev/null +++ b/examples/learning/differentiable_fk/learn_pose_reaching.py @@ -0,0 +1,507 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Learn local Franka reaching through fixed differentiable forward kinematics. + +Run: embodichain run-example learning/differentiable_fk/learn_pose_reaching +Only target TCP poses supervise the network. The example uses the common SL +trainer and does not call an IK solver or start a simulation. +""" + +from __future__ import annotations + +import argparse +from copy import deepcopy +import json +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from torch import nn + +from embodichain.compute.kinematics import SerialChainKinematics, pose_error +from embodichain.data import get_data_path +from embodichain.learning.layers import DifferentiableFKLayer +from embodichain.learning.models import MLP +from embodichain.learning.sl import ( + LossResult, + SupervisedBatch, + SupervisedLearner, + SupervisedTrainer, + SupervisedTrainerCfg, + TensorBatchSource, +) +from embodichain.utils import configclass + +__all__ = ["PoseReachingCfg", "PosePredictor", "load_predictor", "run"] + + +@configclass +class PoseReachingCfg: + """Small, fixed local-reaching experiment; smoke runs check execution only.""" + + device: str = "cpu" + seed: int = 17 + train_count: int = 1024 + validation_count: int = 128 + test_count: int = 128 + updates: int = 1000 + batch_size: int = 128 + learning_rate: float = 1e-3 + state_radius: float = 0.10 + target_radius: float = 0.10 + correction_bound: float = 0.15 + position_scale: float = 0.05 + rotation_scale: float = 0.20 + posture_weight: float = 1e-3 + output: str = "outputs/learn_fk_reaching" + smoke: bool = False + + def __post_init__(self) -> None: + for name in ( + "train_count", + "validation_count", + "test_count", + "updates", + "batch_size", + ): + if type(getattr(self, name)) is not int or getattr(self, name) < 1: + raise ValueError(f"{name} must be a positive integer") + if type(self.seed) is not int or self.seed < 0: + raise ValueError("seed must be a nonnegative integer") + for name in ( + "learning_rate", + "state_radius", + "target_radius", + "correction_bound", + "position_scale", + "rotation_scale", + "posture_weight", + ): + if not np.isfinite(getattr(self, name)) or getattr(self, name) <= 0: + raise ValueError(f"{name} must be finite and positive") + if self.target_radius > self.correction_bound: + raise ValueError("Target perturbations must fit within the output bound") + + +class PosePredictor(nn.Module): + """Example-local joint predictor with a fixed FK layer and bounded corrections. + + Args: + cfg: Feature scales and joint correction bound. + """ + + def __init__(self, cfg: PoseReachingCfg) -> None: + super().__init__() + self.cfg = deepcopy(cfg) + urdf = Path(get_data_path("Franka/Panda/PandaWithHand.urdf")) + self.fk = DifferentiableFKLayer( + SerialChainKinematics.from_urdf(urdf.read_text(), "fr3_hand_tcp", "base") + ) + self.register_buffer( + "ready", torch.tensor([0.0, -0.569, 0.0, -2.810, 0.0, 3.037, 0.741]) + ) + self.network = MLP(13, 7, [64, 64], activation="silu") + nn.init.zeros_(self.network[-1].weight) + nn.init.zeros_(self.network[-1].bias) + + def forward( + self, q_current: torch.Tensor, target_pose: torch.Tensor + ) -> torch.Tensor: + """Predict joints using current joints and a root-frame TCP command. + + Args: + q_current: Admitted current joints, shape ``(..., 7)``. + target_pose: Desired TCP transforms, shape ``(..., 4, 4)``. + + Returns: + Joints within ``correction_bound`` of the current posture. Dataset + sampling reserves that margin inside the physical limits. + """ + position, rotation = pose_error(self.fk(q_current), target_pose) + features = torch.cat( + ( + (q_current - self.ready) / self.cfg.state_radius, + position / self.cfg.position_scale, + rotation / self.cfg.rotation_scale, + ), + -1, + ) + return q_current + self.cfg.correction_bound * self.network(features).tanh() + + +class _PoseObjective: + def compute_loss(self, model: PosePredictor, batch: SupervisedBatch) -> LossResult: + q_pred = model(**batch.inputs) + position, rotation = pose_error(model.fk(q_pred), batch.targets) + pose_loss = ( + (position / model.cfg.position_scale).square().sum(-1) + + (rotation / model.cfg.rotation_scale).square().sum(-1) + ).mean() + posture = (q_pred - batch.inputs["q_current"]).square().sum(-1).mean() + return LossResult( + pose_loss + model.cfg.posture_weight * posture, + {"pose_loss": pose_loss.detach(), "posture_loss": posture.detach()}, + weight=q_pred.shape[0], + ) + + +def _sample(model: PosePredictor, count: int, seed: int) -> SupervisedBatch: + """Generate reachable TCP labels; generating joint solutions are discarded.""" + cfg = model.cfg + low = torch.maximum( + model.ready - cfg.state_radius, model.fk.lower + cfg.correction_bound + ) + high = torch.minimum( + model.ready + cfg.state_radius, model.fk.upper - cfg.correction_bound + ) + if (low > high).any(): + raise ValueError("Reference posture has no admitted sampling interval") + generator = torch.Generator(device="cpu").manual_seed(seed) + q_current = low + torch.rand(count, 7, generator=generator) * (high - low) + q_goal = q_current + cfg.target_radius * ( + 2 * torch.rand(count, 7, generator=generator) - 1 + ) + with torch.no_grad(): + target = model.fk(q_goal) + return SupervisedBatch( + {"q_current": q_current, "target_pose": target}, + target, + metadata={ + "root": "base", + "tcp": "fr3_hand_tcp", + "position_unit": "m", + "joint_unit": "rad", + }, + ) + + +def _errors( + model: PosePredictor, joints: torch.Tensor, targets: torch.Tensor +) -> dict[str, torch.Tensor]: + position, rotation = pose_error(model.fk(joints), targets) + if ( + not torch.isfinite(joints).all() + or not torch.isfinite(position).all() + or not torch.isfinite(rotation).all() + ): + raise FloatingPointError("Non-finite reaching prediction") + if ((joints < model.fk.lower - 1e-6) | (joints > model.fk.upper + 1e-6)).any(): + raise ValueError("Predicted joints exceed physical limits") + p_mm = position.norm(dim=-1) * 1000 + r_deg = torch.rad2deg(rotation.norm(dim=-1)) + return { + "position_mm": p_mm, + "rotation_deg": r_deg, + "success": (p_mm <= 2) & (r_deg <= 2), + } + + +def _summary(errors: dict[str, torch.Tensor]) -> dict[str, float]: + return { + "position_mean_mm": float(errors["position_mm"].mean()), + "position_p95_mm": float(torch.quantile(errors["position_mm"], 0.95)), + "rotation_mean_deg": float(errors["rotation_deg"].mean()), + "rotation_p95_deg": float(torch.quantile(errors["rotation_deg"], 0.95)), + "success_rate": float(errors["success"].float().mean()), + } + + +def _cpu_state(model: nn.Module) -> dict[str, Any]: + return { + name: ( + value.detach().cpu().clone() + if isinstance(value, torch.Tensor) + else deepcopy(value) + ) + for name, value in model.state_dict().items() + } + + +def load_predictor(run_dir: Path, device: str = "cpu") -> PosePredictor: + """Restore a validation-selected frozen predictor on the requested device. + + Args: + run_dir: Directory containing ``predictor.pt``. + device: Inference device, independent of the training device. + + Returns: + An evaluation-mode model with fixed weights and checked FK identity. + """ + saved = torch.load(run_dir / "predictor.pt", map_location="cpu", weights_only=True) + model = PosePredictor(PoseReachingCfg(**saved["config"])) + model.load_state_dict(saved["state_dict"]) + return model.to(device).eval().requires_grad_(False) + + +def _plot(output: Path, history: list[dict[str, float]], best_step: int) -> None: + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + figure, axes = plt.subplots(1, 2, figsize=(10, 3.8)) + for axis, key, label in zip( + axes, + ("position_mean_mm", "rotation_mean_deg"), + ("Position error / mm", "Rotation error / degrees"), + ): + axis.plot( + [row["step"] for row in history], + [row[key] for row in history], + label="Independent validation", + ) + axis.axvline( + best_step, color="tab:green", linestyle="--", label="Selected checkpoint" + ) + axis.set(xlabel="Optimizer update", ylabel=label) + axis.grid(alpha=0.2) + axis.legend(fontsize=8) + figure.tight_layout() + figure.savefig(output / "learning.png", dpi=160) + plt.close(figure) + + +def run(cfg: PoseReachingCfg) -> dict[str, Any]: + """Train with pose supervision and evaluate every independent test sample. + + Args: + cfg: Experiment sizes, numerical settings and output location. + + Returns: + Metrics, selected update, learning checks and configuration. + """ + cfg.validate() + torch.set_num_threads(4) + torch.manual_seed(cfg.seed) + model = PosePredictor(cfg) + datasets = { + name: _sample(model, count, seed) + for name, count, seed in ( + ("train", cfg.train_count, 23), + ("validation", cfg.validation_count, 24), + ("test", cfg.test_count, 25), + ) + } + source = TensorBatchSource(datasets["train"], source_id="franka-fk-reaching-v1") + model.to(cfg.device) + fixed_fk = _cpu_state(model.fk) + objective = _PoseObjective() + learner = SupervisedLearner( + model, objective, torch.optim.Adam(model.parameters(), lr=cfg.learning_rate) + ) + validation = datasets["validation"].to(cfg.device) + history: list[dict[str, float]] = [] + best_loss, best_step, best_state = float("inf"), 0, _cpu_state(model) + + def evaluate(current: SupervisedLearner) -> dict[str, float]: + nonlocal best_loss, best_step, best_state + with torch.no_grad(): + joints = model(**validation.inputs) + stats = _summary(_errors(model, joints, validation.targets)) + value = float(objective.compute_loss(model, validation).loss) + history.append({"step": current.num_updates, "loss": value, **stats}) + if value < best_loss: + best_loss, best_step, best_state = ( + value, + current.num_updates, + _cpu_state(model), + ) + print( + f"step={current.num_updates:4d} validation={value:.6g} position={stats['position_mean_mm']:.3f} mm rotation={stats['rotation_mean_deg']:.3f} deg", + flush=True, + ) + return {"loss": value, **stats} + + evaluate(learner) + interval = min(50, cfg.updates) + trainer = SupervisedTrainer( + learner, + source, + SupervisedTrainerCfg( + total_updates=cfg.updates, + batch_size=cfg.batch_size, + device=cfg.device, + seed=cfg.seed + 12, + output_dir=cfg.output, + checkpoint_interval=0, + validation_interval=interval, + ), + signature={ + "recipe": "franka-fk-reaching-v1", + "config": cfg.to_dict(), + "model": "13-64-64-7-silu", + "selection": "validation-loss-only", + }, + evaluate=evaluate, + ) + trainer.fit() + if cfg.updates % interval: + evaluate(learner) + trained_fk = _cpu_state(model.fk) + for name, value in fixed_fk.items(): + if isinstance(value, torch.Tensor): + torch.testing.assert_close(value, trained_fk[name], atol=0, rtol=0) + else: + assert value == trained_fk[name] + model.load_state_dict(best_state) + test = datasets["test"].to(cfg.device) + with torch.no_grad(): + before_q = test.inputs["q_current"] + after_q = model(**test.inputs) + before = _errors(model, before_q, test.targets) + after = _errors(model, after_q, test.targets) + # Check only the pose term: the posture regularizer cannot hide a detached FK. + probe = SupervisedBatch( + {name: value[:4] for name, value in validation.inputs.items()}, + validation.targets[:4], + ) + predicted = model(**probe.inputs) + p, r = pose_error(model.fk(predicted), probe.targets) + pose_loss = (p.square().sum(-1) + r.square().sum(-1)).mean() + gradients = torch.autograd.grad(pose_loss, tuple(model.network.parameters())) + assert all(torch.isfinite(g).all() for g in gradients) + assert sum(float(g.abs().sum()) for g in gradients) > 0 + + output = Path(cfg.output) + torch.save( + {"config": cfg.to_dict(), "state_dict": best_state}, output / "predictor.pt" + ) + reloaded = load_predictor(output) + with torch.no_grad(): + restored_q = reloaded(**datasets["test"].inputs) + torch.testing.assert_close(restored_q, after_q.cpu(), atol=2e-5, rtol=1e-4) + np.savez_compressed( + output / "datasets.npz", + **{ + f"{split}_{name}": tensor.numpy() + for split, batch in datasets.items() + for name, tensor in batch.inputs.items() + }, + ) + np.savez_compressed( + output / "evaluation.npz", + current_q=before_q.cpu().numpy(), + target_pose=test.targets.cpu().numpy(), + after_q=after_q.cpu().numpy(), + **{ + f"{method}_{name}": value.cpu().numpy() + for method, errors in (("before", before), ("after", after)) + for name, value in errors.items() + }, + ) + before_stats, after_stats = _summary(before), _summary(after) + improvement = { + name: 1 - after_stats[key] / before_stats[key] + for name, key in ( + ("position", "position_mean_mm"), + ("rotation", "rotation_mean_deg"), + ) + } + passed = all(value >= 0.5 for value in improvement.values()) + metrics = { + "config": cfg.to_dict(), + "best_step": best_step, + "history": history, + "test_count": cfg.test_count, + "before": before_stats, + "after": after_stats, + "improvement": improvement, + "learning_acceptance_passed": passed, + "checks": { + "pose_gradients": True, + "fixed_fk": True, + "physical_limits": True, + "cpu_reload": True, + }, + } + (output / "metrics.json").write_text( + json.dumps(metrics, indent=2, allow_nan=False) + "\n" + ) + _plot(output, history, best_step) + rows = [ + "# Learning reaching through differentiable FK", + "", + f"Independent test samples: {cfg.test_count}; validation-selected update: {best_step}.", + "", + "| Model | Position mean / mm | Position P95 / mm | Rotation mean / deg | Rotation P95 / deg | Within 2 mm and 2 deg |", + "| --- | --- | --- | --- | --- | --- |", + ] + for name, stats in ( + ("Before training", before_stats), + ("After training", after_stats), + ): + rows.append( + f"| {name} | {stats['position_mean_mm']:.4f} | {stats['position_p95_mm']:.4f} | {stats['rotation_mean_deg']:.4f} | {stats['rotation_p95_deg']:.4f} | {stats['success_rate']:.1%} |" + ) + rows.extend( + [ + "", + "![Validation curves](learning.png)", + "", + "All test samples, including failures, contribute to these metrics. Targets are local and reachable by construction. Smoke runs check execution only. The viewer shows kinematic configurations without physical execution.", + ] + ) + (output / "report.md").write_text("\n".join(rows) + "\n") + print( + json.dumps( + { + "output": str(output), + "before": before_stats, + "after": after_stats, + "improvement": improvement, + }, + indent=2, + ), + flush=True, + ) + if not cfg.smoke and not passed: + raise RuntimeError( + "Held-out position and rotation mean errors must each improve by at least 50%; inspect report.md" + ) + return metrics + + +def main() -> None: + """Train and evaluate a standalone FK-supervised joint predictor.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--device", default="cpu") + parser.add_argument("--seed", type=int, default=17) + parser.add_argument("--updates", type=int, default=1000) + parser.add_argument("--output", default="outputs/learn_fk_reaching") + parser.add_argument("--smoke", action="store_true") + args = parser.parse_args() + options = dict( + device=args.device, + seed=args.seed, + updates=args.updates, + output=args.output, + smoke=args.smoke, + ) + if args.smoke: + options.update( + updates=20, + train_count=64, + validation_count=16, + test_count=16, + batch_size=32, + ) + run(PoseReachingCfg(**options)) + + +if __name__ == "__main__": + main() diff --git a/examples/learning/differentiable_fk/requirements.txt b/examples/learning/differentiable_fk/requirements.txt new file mode 100644 index 000000000..365679623 --- /dev/null +++ b/examples/learning/differentiable_fk/requirements.txt @@ -0,0 +1,5 @@ +# Optional plotting and mesh-viewer dependencies for this example. +matplotlib +plotly +trimesh +yourdfpy diff --git a/examples/learning/differentiable_fk/view_pose_reaching.py b/examples/learning/differentiable_fk/view_pose_reaching.py new file mode 100644 index 000000000..2cab5a7ac --- /dev/null +++ b/examples/learning/differentiable_fk/view_pose_reaching.py @@ -0,0 +1,331 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Preview a frozen FK-supervised Franka predictor against its initial posture. + +Run: embodichain run-example learning/differentiable_fk/view_pose_reaching +Train with learn_pose_reaching first. This viewer performs CPU inference only. +""" + +from __future__ import annotations + +import argparse +from concurrent.futures import Future, ThreadPoolExecutor +import json +from pathlib import Path +import time +from typing import Any +import urllib.request + +import numpy as np +import plotly.graph_objects as go +from plotly.subplots import make_subplots +import torch +import viser +from viser.extras import ViserUrdf +from viser.transforms import SO3 + +from embodichain.compute.kinematics import pose_error +from embodichain.data import get_data_path +from .learn_pose_reaching import _cpu_state, load_predictor + +__all__ = ["serve"] + + +def _learning_curve(history: list[dict[str, float]]) -> go.Figure: + figure = make_subplots(rows=2, cols=1, shared_xaxes=True, vertical_spacing=0.18) + for row, name, label in ( + (1, "position_mean_mm", "Position / mm"), + (2, "rotation_mean_deg", "Rotation / deg"), + ): + figure.add_trace( + go.Scatter( + x=[item["step"] for item in history], + y=[item[name] for item in history], + mode="lines+markers", + showlegend=False, + ), + row=row, + col=1, + ) + figure.update_yaxes(title_text=label, row=row, col=1) + figure.update_xaxes(title_text="Training update", row=2, col=1) + figure.update_layout( + title="Independent validation", + height=330, + margin=dict(l=55, r=10, t=35, b=35), + font=dict(size=10), + ) + return figure + + +def serve( + run_dir: Path, host: str = "127.0.0.1", port: int = 8103, smoke: bool = False +) -> None: + """Compare held-out and edited TCP targets using a frozen CPU predictor. + + Args: + run_dir: Training directory containing exported weights and evaluation. + host: Local server bind address. + port: Preferred HTTP port; Viser reports the actual selected port. + smoke: Check scene construction, one target edit and HTTP, then stop. + """ + torch.set_num_threads(4) + metadata = json.loads((run_dir / "metrics.json").read_text()) + with np.load(run_dir / "evaluation.npz", allow_pickle=False) as saved: + current_q, targets = saved["current_q"].copy(), saved["target_pose"].copy() + model = load_predictor(run_dir) + frozen = _cpu_state(model) + urdf = Path(get_data_path("Franka/Panda/PandaWithHand.urdf")) + server = viser.ViserServer( + host=host, port=port, label="Learn reaching through differentiable FK" + ) + server.scene.set_up_direction("+z") + server.scene.add_grid( + "/floor", width=2.5, height=3.2, cell_size=0.1, section_size=0.5 + ) + roots, robots, actual_frames, target_frames, lines, goal_labels = ( + [], + [], + [], + [], + [], + [], + ) + for name, offset, color, title in ( + ("before", -0.65, (0.34, 0.55, 0.91), "Before training"), + ("after", 0.65, (0.13, 0.68, 0.54), "After training"), + ): + root = "/" + name + roots.append( + server.scene.add_frame(root, show_axes=False, position=(0.0, offset, 0.0)) + ) + robots.append( + ViserUrdf(server, urdf, root_node_name=root, mesh_color_override=color) + ) + actual_frames.append( + server.scene.add_frame(root + "/tcp", axes_length=0.06, axes_radius=0.002) + ) + target_frames.append( + server.scene.add_frame( + root + "/target", axes_length=0.10, axes_radius=0.003 + ) + ) + lines.append( + server.scene.add_line_segments( + root + "/error", + points=np.zeros((1, 2, 3), dtype=np.float32), + colors=(244, 172, 55), + line_width=4, + ) + ) + goal_labels.append(server.scene.add_label(root + "/goal_label", "Goal TCP")) + server.scene.add_label(root + "/title", title, position=(0.0, 0.0, 1.12)) + + @server.on_client_connect + def camera(client: viser.ClientHandle) -> None: + client.camera.position = (2.4, 0.0, 1.5) + client.camera.look_at = (0.25, 0.0, 0.45) + client.camera.up_direction = (0.0, 0.0, 1.0) + client.camera.fov = float(np.deg2rad(45)) + + server.gui.add_markdown( + "## Reaching learned through FK\nA network predicts joints from **current joints + target pose**. Blue: zero-initialized output. Green: validation-selected model. Large axes mark the goal; yellow lines show position error." + ) + case = server.gui.add_slider( + "Test case", 1, max(2, len(current_q)), 1, 1, disabled=len(current_q) == 1 + ) + with server.gui.add_folder("Edit target", expand_by_default=True): + translation = [ + server.gui.add_slider(f"Offset {axis} (m)", -0.05, 0.05, 0.001, 0.0) + for axis in "XYZ" + ] + angles = [ + server.gui.add_slider(f"{name} (deg)", -15.0, 15.0, 1.0, 0.0) + for name in ("Roll", "Pitch", "Yaw") + ] + reset = server.gui.add_button("Reset offsets") + controls = translation + angles + + @reset.on_click + def reset_offsets(_: viser.GuiEvent) -> None: + for control in controls: + control.value = 0.0 + + status = server.gui.add_markdown("Loading prediction…") + errors = server.gui.add_markdown("") + server.gui.add_plotly(_learning_curve(metadata["history"]), aspect=1.1) + summary = [ + f"### All {metadata['test_count']} held-out targets", + "| Metric | Before | After |", + "| --- | --- | --- |", + ] + for key, label in ( + ("position_mean_mm", "Position mean / mm"), + ("position_p95_mm", "Position P95 / mm"), + ("rotation_mean_deg", "Rotation mean / deg"), + ("rotation_p95_deg", "Rotation P95 / deg"), + ): + summary.append( + f"| {label} | {metadata['before'][key]:.3f} | {metadata['after'][key]:.3f} |" + ) + summary.append( + f"| Within 2 mm and 2 deg | {metadata['before']['success_rate']:.1%} | {metadata['after']['success_rate']:.1%} |" + ) + server.gui.add_markdown("\n".join(summary)) + server.gui.add_markdown( + "Weights stay frozen. Edited goals are outside the saved evaluation and may exceed the learned local region. This is a kinematic preview without collision, contact or physical execution." + ) + + def infer(snapshot: tuple[float, ...]) -> dict[str, Any]: + index = int(snapshot[0]) - 1 + target = targets[index].copy() + target[:3, 3] += np.asarray(snapshot[1:4], dtype=np.float32) + rotation = SO3.from_rpy_radians(*np.deg2rad(snapshot[4:7])).as_matrix() + target[:3, :3] = rotation @ target[:3, :3] + q = torch.from_numpy(current_q[index : index + 1].copy()) + goal = torch.from_numpy(target[None]) + with torch.no_grad(): + predicted = model(q, goal) + joints = torch.cat((q, predicted)) + poses = model.fk(joints) + position, orientation = pose_error(poses, goal) + if not torch.isfinite(joints).all(): + raise FloatingPointError("Non-finite frozen prediction") + if ((joints < model.fk.lower - 1e-6) | (joints > model.fk.upper + 1e-6)).any(): + raise ValueError("Frozen prediction exceeds joint limits") + return { + "q": joints.numpy(), + "poses": poses.numpy(), + "target": target, + "position_mm": position.norm(dim=-1).numpy() * 1000, + "rotation_deg": torch.rad2deg(orientation.norm(dim=-1)).numpy(), + "edited": any(value != 0 for value in snapshot[1:]), + "case": index + 1, + } + + def render(result: dict[str, Any]) -> None: + for j, robot in enumerate(robots): + lookup = dict(zip(model.fk.model.joint_names, result["q"][j])) + robot.update_cfg( + np.asarray( + [ + lookup.get(name, 0.04) + for name in robot.get_actuated_joint_names() + ] + ) + ) + actual_frames[j].position = result["poses"][j, :3, 3] + actual_frames[j].wxyz = SO3.from_matrix(result["poses"][j, :3, :3]).wxyz + target_frames[j].position = result["target"][:3, 3] + target_frames[j].wxyz = SO3.from_matrix(result["target"][:3, :3]).wxyz + lines[j].points = np.stack( + (result["poses"][j, :3, 3], result["target"][:3, 3]) + )[None] + goal_labels[j].position = result["target"][:3, 3] + np.array( + [0.0, 0.0, 0.07] + ) + roots[j].visible = True + status.content = f"**{'Edited target · outside saved evaluation' if result['edited'] else 'Held-out test target'} {result['case']}**\n\nFrozen checkpoint: update {metadata['best_step']}." + rows = ["| Current target | Before | After |", "| --- | --- | --- |"] + for key, label in ( + ("position_mm", "Position / mm"), + ("rotation_deg", "Rotation / deg"), + ): + rows.append(f"| {label} | {result[key][0]:.3f} | {result[key][1]:.3f} |") + errors.content = "\n".join(rows) + + executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="fk-preview") + future: Future[dict[str, Any]] | None = None + pending: tuple[int, tuple[float, ...]] | None = None + requested: tuple[float, ...] | None = None + revision = active_revision = 0 + previous_case = int(case.value) + smoke_edited = False + print(f"Preview: http://{host}:{server.get_port()}", flush=True) + try: + while True: + if int(case.value) != previous_case: + for control in controls: + control.value = 0.0 + previous_case = int(case.value) + snapshot = ( + float(case.value), + *(float(control.value) for control in controls), + ) + if snapshot != requested: + requested, revision = snapshot, revision + 1 + pending = (revision, snapshot) + status.content = "Computing frozen prediction…" + errors.content = "" + for root in roots: + root.visible = False + if future is not None and future.done(): + result = future.result() + future = None + if active_revision == revision: + render(result) + if smoke and not smoke_edited: + translation[0].value = 0.01 + smoke_edited = True + elif smoke: + assert result["edited"] + assert all(robot.get_actuated_joint_names() for robot in robots) + assert not any( + parameter.requires_grad for parameter in model.parameters() + ) + for name, value in _cpu_state(model).items(): + if isinstance(value, torch.Tensor): + torch.testing.assert_close( + value, frozen[name], atol=0, rtol=0 + ) + else: + assert value == frozen[name] + connect_host = "127.0.0.1" if host == "0.0.0.0" else host + with urllib.request.urlopen( + f"http://{connect_host}:{server.get_port()}", timeout=10 + ) as response: + assert response.status == 200 + print( + "Viewer smoke passed: robot scenes, target edit, bounded predictions, frozen weights and HTTP", + flush=True, + ) + break + if future is None and pending is not None: + active_revision, work = pending + pending = None + future = executor.submit(infer, work) + time.sleep(0.02) + except KeyboardInterrupt: + pass + finally: + executor.shutdown(wait=True, cancel_futures=True) + server.stop() + + +def main() -> None: + """Open the frozen-model comparison, or run its bounded smoke check.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--run", type=Path, default=Path("outputs/learn_fk_reaching")) + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, default=8103) + parser.add_argument("--smoke", action="store_true") + args = parser.parse_args() + serve(args.run, args.host, args.port, args.smoke) + + +if __name__ == "__main__": + main() diff --git a/examples/learning/differentiable_ik/README.md b/examples/learning/differentiable_ik/README.md index f6c345ee5..cf8204a6e 100644 --- a/examples/learning/differentiable_ik/README.md +++ b/examples/learning/differentiable_ik/README.md @@ -6,6 +6,12 @@ grasp position**. The position and downward approach remain fixed. During traini plus pose residuals backpropagate into the predictor. After training, frozen weights handle new tasks with one network evaluation and one fixed-step IK solve. +The shared `DifferentiableFKLayer` maps joints to TCP poses with `fk(qpos)`; +`DifferentiableIKLayer(fk, cfg)` composes it to solve the reverse mapping. FK also +supports pose supervision for a network that predicts joints: apply `pose_error` +to the predicted and target TCP transforms and backpropagate that loss. Geometry +and TCP stay fixed, and task code owns loss weights and joint-limit handling. + ```text (q_current, grasp_position) → MLP → yaw → TCP target → differentiable IK ↑ ↓ diff --git a/examples/learning/differentiable_ik/learn_grasp_yaw.py b/examples/learning/differentiable_ik/learn_grasp_yaw.py index 7f52dee80..3e8b9efc0 100644 --- a/examples/learning/differentiable_ik/learn_grasp_yaw.py +++ b/examples/learning/differentiable_ik/learn_grasp_yaw.py @@ -38,7 +38,7 @@ from embodichain.compute.kinematics import DifferentiableIKCfg, SerialChainKinematics from embodichain.data import get_data_path -from embodichain.learning.layers import DifferentiableIKLayer, KinematicsLayer +from embodichain.learning.layers import DifferentiableIKLayer, DifferentiableFKLayer from embodichain.learning.models import MLP from embodichain.utils import configclass from ._benchmark import _measure, _sync, _table @@ -90,10 +90,10 @@ def __post_init__(self) -> None: raise ValueError("learning_rate must be finite and positive") -def load_kinematics(device: str = "cpu") -> KinematicsLayer: +def load_kinematics(device: str = "cpu") -> DifferentiableFKLayer: """Build the official Franka arm chain without importing simulation.""" path = Path(get_data_path("Franka/Panda/PandaWithHand.urdf")) - return KinematicsLayer( + return DifferentiableFKLayer( SerialChainKinematics.from_urdf( path.read_text(), "fr3_hand_tcp", "base", device=device ) @@ -116,12 +116,12 @@ def target_pose(angle: torch.Tensor, position: torch.Tensor) -> torch.Tensor: ) -def ready_joints(k: KinematicsLayer) -> torch.Tensor: +def ready_joints(k: DifferentiableFKLayer) -> torch.Tensor: """Return the official reference posture in the model's dtype/device.""" return k.tcp.new_tensor([0.0, -0.569, 0.0, -2.810, 0.0, 3.037, 0.741]) -def current_heading(k: KinematicsLayer, q: torch.Tensor) -> torch.Tensor: +def current_heading(k: DifferentiableFKLayer, q: torch.Tensor) -> torch.Tensor: """Project current TCP X into root XY; data stays near downward grasps.""" pose = k(q) return torch.atan2(pose[..., 1, 0], pose[..., 0, 0]) @@ -134,7 +134,7 @@ class GraspYawPredictor(nn.Module): declared task scales, not statistics from validation or test targets. """ - def __init__(self, k: KinematicsLayer) -> None: + def __init__(self, k: DifferentiableFKLayer) -> None: super().__init__() base = ready_joints(k) with torch.no_grad(): @@ -161,7 +161,7 @@ def forward( class ConstantYaw(nn.Module): """Fit one absolute grasp angle shared by every task, without conditioning.""" - def __init__(self, k: KinematicsLayer) -> None: + def __init__(self, k: DifferentiableFKLayer) -> None: super().__init__() self.angle = nn.Parameter(k.tcp.new_zeros(())) @@ -174,7 +174,7 @@ def forward( @torch.no_grad() def sample_tasks( - k: KinematicsLayer, count: int, seed: int, split: str + k: DifferentiableFKLayer, count: int, seed: int, split: str ) -> dict[str, torch.Tensor]: """Sample independent local problems with a disjoint Cartesian-Y holdout.""" g = torch.Generator().manual_seed(seed) @@ -313,7 +313,7 @@ def grid_reference( def fit_model( - k: KinematicsLayer, + k: DifferentiableFKLayer, cfg: GraspLearningCfg, train: dict[str, torch.Tensor], validation: dict[str, torch.Tensor], @@ -432,7 +432,9 @@ def summarize( } -def validate_gradients(k: KinematicsLayer, iterations: int = 24) -> dict[str, Any]: +def validate_gradients( + k: DifferentiableFKLayer, iterations: int = 24 +) -> dict[str, Any]: """Check complete policy→IK task gradients and reject hidden input truncation.""" # These are executable demo checks, not a new tutorial-importing pytest suite. torch.manual_seed(101) @@ -482,7 +484,7 @@ def loss_fn(q_value: torch.Tensor, p_value: torch.Tensor) -> torch.Tensor: def load_predictors( run_dir: Path, device: str = "cpu" -) -> tuple[KinematicsLayer, dict[str, nn.Module], GraspLearningCfg]: +) -> tuple[DifferentiableFKLayer, dict[str, nn.Module], GraspLearningCfg]: """Restore the first seeded run's validation-selected, frozen predictors.""" meta = json.loads((run_dir / "metrics.json").read_text()) cfg = GraspLearningCfg(**meta["config"]) diff --git a/scripts/benchmark/robotics/kinematic_solver/run_benchmark.py b/scripts/benchmark/robotics/kinematic_solver/run_benchmark.py index a016e865d..89afbce64 100644 --- a/scripts/benchmark/robotics/kinematic_solver/run_benchmark.py +++ b/scripts/benchmark/robotics/kinematic_solver/run_benchmark.py @@ -26,6 +26,7 @@ from __future__ import annotations import argparse +from collections.abc import Callable from importlib.metadata import version import os import time @@ -86,11 +87,11 @@ def _parse_args() -> argparse.Namespace: "--solvers", "-s", nargs="+", - choices=(*SUPPORTED_SOLVERS, "franka", "ssik", "all"), + choices=(*SUPPORTED_SOLVERS, "franka", "ssik", "fk", "all"), default=["all"], help=( "Solvers to benchmark. Use one or more of: opw, pytorch, ur, fep, differentiable, " - "franka, ssik, all. Franka and SSIK comparisons run only when selected explicitly. " + "franka, ssik, fk, all. Franka, SSIK and FK comparisons run only when selected explicitly. " "Default: all." ), ) @@ -221,13 +222,13 @@ def _normalize_selected_solvers(selected_solvers: list[str] | None) -> set[str]: """Normalize selected solver names to a canonical set.""" if not selected_solvers or "all" in selected_solvers: explicit_comparisons = set(selected_solvers or ()).intersection( - ("franka", "ssik") + ("franka", "ssik", "fk") ) return set(SUPPORTED_SOLVERS) | explicit_comparisons return { solver for solver in selected_solvers - if solver in (*SUPPORTED_SOLVERS, "franka", "ssik") + if solver in (*SUPPORTED_SOLVERS, "franka", "ssik", "fk") } @@ -1327,7 +1328,7 @@ def benchmark_differentiable_solver() -> ( damped_least_squares_step, pose_error, ) - from embodichain.learning.layers import DifferentiableIKLayer, KinematicsLayer + from embodichain.learning.layers import DifferentiableIKLayer, DifferentiableFKLayer perf_rows: list[dict[str, object]] = [] metric_rows: list[dict[str, object]] = [] @@ -1395,7 +1396,7 @@ def measure(call, device): for device in devices: model = cpu_model.to(device, torch.float32) - kinematics = KinematicsLayer(model) + kinematics = DifferentiableFKLayer(model) common_cfg = dict( urdf_path=path, root_link_name="base", @@ -1554,6 +1555,155 @@ def call_baseline(): return perf_rows, metric_rows +def benchmark_differentiable_fk() -> ( + tuple[list[dict[str, object]], list[dict[str, object]]] +): + """Compare pure FK with pose extraction through FK/Jacobian evaluation. + + Use identical seven-DOF Franka inputs, a translated/rotated TCP, and pose + losses for both implementations. Time device-resident steady-state calls; + measure warmed memory separately, retaining one output until after sampling. + """ + from embodichain.compute.kinematics import SerialChainKinematics, pose_error + from embodichain.learning.layers import DifferentiableFKLayer + + model = SerialChainKinematics.from_urdf( + Path(get_data_path("Franka/Panda/PandaWithHand.urdf")).read_text(), + "fr3_hand_tcp", + "base", + ) + tcp = torch.tensor( + [[0, -1, 0, 0.1], [1, 0, 0, -0.05], [0, 0, 1, 0.08], [0, 0, 0, 1]], + dtype=torch.float32, + ) + generator = torch.Generator().manual_seed(2026) + workloads = [] + for count in (1, 64, 1024): + q = model.chain.low + ( + 0.25 + 0.5 * torch.rand(count, model.dof, generator=generator) + ) * (model.chain.high - model.chain.low) + target = model.forward_kinematics(q + 0.04, tcp) + workloads.append((q, target)) + + def measure( + call: Callable[[], torch.Tensor], device: torch.device + ) -> dict[str, object]: + def synchronize() -> None: + if device.type == "cuda": + torch.cuda.synchronize(device) + + for _ in range(5): + output = call() + del output + times = [] + for _ in range(20): + synchronize() + start = time.perf_counter() + output = call() + synchronize() + times.append((time.perf_counter() - start) * 1000) + del output + + # RSS measures retained memory, not CPU transient peak. Torch's CUDA + # counters cover allocations by these Torch/PK paths, not total VRAM. + synchronize() + process = psutil.Process() + cpu_before = process.memory_info().rss + cuda = device.type == "cuda" + gpu_before = torch.cuda.memory_allocated(device) if cuda else 0 + if cuda: + torch.cuda.reset_peak_memory_stats(device) + output = call() + synchronize() + cpu_delta = (process.memory_info().rss - cpu_before) / 1024**2 + gpu_delta = ( + (torch.cuda.memory_allocated(device) - gpu_before) / 1024**2 + if cuda + else "N/A" + ) + peak = torch.cuda.max_memory_allocated(device) if cuda else 0 + del output + return { + "cost_time_ms": float(np.median(times)), + "p90_time_ms": float(np.percentile(times, 90)), + "cpu_delta_mb": cpu_delta, + "gpu_delta_mb": gpu_delta, + "peak_gpu_mb": peak / 1024**2 if cuda else "N/A", + "peak_gpu_delta_mb": (peak - gpu_before) / 1024**2 if cuda else "N/A", + } + + devices = [torch.device("cpu")] + if torch.cuda.is_available(): + devices.append(torch.device("cuda", torch.cuda.current_device())) + perf_rows: list[dict[str, object]] = [] + metric_rows: list[dict[str, object]] = [] + for device in devices: + fk = DifferentiableFKLayer(model, tcp).to(device) + for cpu_q, cpu_target in workloads: + q = cpu_q.to(device).detach().requires_grad_() + target = cpu_target.to(device) + + def loss(pose: torch.Tensor) -> torch.Tensor: + position, rotation = pose_error(pose, target) + return (position.square().sum(-1) + rotation.square().sum(-1)).mean() + + reference = fk.evaluate(q)[0] + reference_gradient = torch.autograd.grad(loss(reference), q)[0] + reference = reference.detach() + for name, operation in ( + ("fk_only", fk), + ("fk_via_jacobian", lambda joints: fk.evaluate(joints)[0]), + ): + pose = operation(q) + gradient = torch.autograd.grad(loss(pose), q)[0] + torch.testing.assert_close(pose, reference) + torch.testing.assert_close(gradient, reference_gradient) + with torch.no_grad(): + position, rotation = pose_error(pose, reference) + position_norm = position.norm(dim=-1) + rotation_norm = rotation.norm(dim=-1) + common = { + "sample_size": q.shape[0], + "impl": f"{name}_{device.type}", + "component": "fk_pose", + } + metric_rows.append( + { + **common, + "success_rate": ( + (position_norm < 1e-6) & (rotation_norm < 1e-6) + ) + .float() + .mean() + .item(), + "translation_err_mm": position_norm.mean().item() * 1000, + "rotation_err_deg": rotation_norm.mean().item() * 180 / np.pi, + "gradient_max_abs_error": (gradient - reference_gradient) + .abs() + .max() + .item(), + } + ) + del pose, gradient + for mode in ("no_grad", "forward", "forward_backward"): + + def call() -> torch.Tensor: + with torch.set_grad_enabled(mode != "no_grad"): + pose = operation(q) + if mode == "forward_backward": + return torch.autograd.grad(loss(pose), q)[0] + return pose + + measured = measure(call, device) + perf_rows.append({**common, "mode": mode, **measured}) + print( + f" {name} {device.type} n={q.shape[0]} {mode}: " + f"{measured['cost_time_ms']:.3f} ms", + flush=True, + ) + return perf_rows, metric_rows + + def run_all_benchmarks( selected_solvers: list[str] | None = None, *, @@ -1584,6 +1734,8 @@ def run_all_benchmarks( print( "- Franka comparison: default FEP search versus former Pytorch with 30 seeds." ) + if "fk" in solvers_to_run: + print("- FK comparison: pure FK versus pose extraction through FK/Jacobian.") perf_rows: list[dict[str, object]] = [] metric_rows: list[dict[str, object]] = [] @@ -1618,6 +1770,11 @@ def run_all_benchmarks( perf_rows.extend(diff_perf) metric_rows.extend(diff_metrics) + if "fk" in solvers_to_run: + fk_perf, fk_metrics = benchmark_differentiable_fk() + perf_rows.extend(fk_perf) + metric_rows.extend(fk_metrics) + ssik_notes: list[str] = [] if "ssik" in solvers_to_run: from scripts.benchmark.robotics.kinematic_solver.ssik_benchmark import ( @@ -1643,9 +1800,16 @@ def run_all_benchmarks( leaderboard_rows=leaderboard_rows, notes=[ *ssik_notes, - "Legacy CPU/GPU memory fields are deltas measured around timed calls; SSIK uses a separate warmed memory pass with separately named Torch and Warp fields.", - f"Runtime: Torch {torch.__version__}, pytorch_kinematics {version('pytorch_kinematics')}, CPU threads {torch.get_num_threads()}, GPU {torch.cuda.get_device_name() if torch.cuda.is_available() else 'unavailable'}. Numeric benchmark inputs use float32.", + "Legacy CPU/GPU memory fields are deltas measured around timed calls; SSIK and FK use separate warmed memory passes. SSIK names Torch and Warp fields separately.", + f"Runtime: Torch {torch.__version__}, CUDA {torch.version.cuda}, pytorch_kinematics {version('pytorch_kinematics')}, CPU threads {torch.get_num_threads()}, GPU {torch.cuda.get_device_name() if torch.cuda.is_available() else 'unavailable'}. Numeric benchmark inputs use float32.", "This report contains exactly three tables: Time & Memory, Success & Other Metrics, and Leaderboard.", + *( + [ + "FK mode: float32 Franka with translated/rotated TCP, batches 1/64/1024 and CPU RNG seed 2026. Inputs and targets are generated once on CPU, then moved before timing. Median and p90 completed-call wall latency use 20 synchronized repetitions after five warmups. no_grad measures inference; forward retains the graph; forward_backward rebuilds the graph and differentiates the same position/rotation loss without accumulating leaf gradients. Memory uses one separate warmed call retaining its output; CPU delta is retained RSS, not transient peak. CUDA peak_gpu_mb is total Torch peak allocated memory; peak_gpu_delta_mb is additional peak above the warmed baseline. CPU-only CUDA metrics are N/A. FK success_rate is pose agreement within 1e-6 m/rad of evaluate()[0]; gradient_max_abs_error compares input gradients. Geometry construction and transfers are excluded." + ] + if "fk" in solvers_to_run + else [] + ), *( [ "Differentiable mode uses shared CPU-generated Franka targets/seeds on CPU and CUDA. DLS: 4/8/16 steps, damping .05, step size .5, hard limits. FEP: redundancy search. Pytorch: two seeds, 64 iterations, lr .5. Timings are medians of three synchronized calls after two warmups, excluding data preparation and RNG reset to seed 2026; each backward call rebuilds its graph. Pose errors include all rows, including failures. Component rows are excluded from the success leaderboard. Peak GPU memory is absolute allocated Torch memory; Warp allocations in FEP are not captured by that counter." diff --git a/tests/compute/kinematics/test_dls.py b/tests/compute/kinematics/test_dls.py index 9748dc26b..c502e5c56 100644 --- a/tests/compute/kinematics/test_dls.py +++ b/tests/compute/kinematics/test_dls.py @@ -101,6 +101,14 @@ def test_tcp_jacobian_forward_and_backward(): assert torch.autograd.gradcheck( lambda joints: evaluate(joints)[1], (q,), fast_mode=True ) + fk = lambda joints: model.forward_kinematics(joints, tcp) + assert torch.autograd.gradcheck(fk, (q,), fast_mode=True) + torch.testing.assert_close(fk(q), evaluate(q)[0]) + weights = torch.arange(16, dtype=q.dtype).reshape(4, 4) + torch.testing.assert_close( + torch.autograd.grad((fk(q) * weights).sum(), q)[0], + torch.autograd.grad((evaluate(q)[0] * weights).sum(), q)[0], + ) reference = evaluate(q)[1] epsilon = 1e-6 columns = [] @@ -115,6 +123,83 @@ def test_tcp_jacobian_forward_and_backward(): ) +def test_forward_kinematics_planar_chain(): + # A unit upper arm and half-unit forearm rotate about parallel Z axes. + # Side branches ensure that the controlled tip need not be the last frame. + urdf = """ + + + + + + + + + + + + + + + + + + + + """ + model = SerialChainKinematics.from_urdf(urdf, "tip", dtype=torch.float64) + tcp = torch.tensor( + [[0, -1, 0, 0.1], [1, 0, 0, 0.2], [0, 0, 1, 0.3], [0, 0, 0, 1]], + dtype=torch.float64, + ) + # Both joints at pi/2 put the flange at (-0.5, 1, 0). Its rotation by pi + # moves the TCP offset to (-0.1, -0.2, 0.3), even beyond physical limits. + q = torch.tensor([[0, 0], [torch.pi / 2, torch.pi / 2]], dtype=torch.float64) + expected = torch.tensor( + [ + [[0, -1, 0, 1.6], [1, 0, 0, 0.2], [0, 0, 1, 0.3], [0, 0, 0, 1]], + [[0, 1, 0, -0.6], [-1, 0, 0, 0.8], [0, 0, 1, 0.3], [0, 0, 0, 1]], + ], + dtype=torch.float64, + ) + torch.testing.assert_close(model.forward_kinematics(q, tcp), expected) + + +@pytest.mark.parametrize("batch", [(), (2, 3), (0,)]) +def test_forward_kinematics_batch_layout_and_per_sample_tcp(batch): + model, tcp = model_and_tcp() + # Move the joint axis to make multidimensional inputs noncontiguous. + q = torch.linspace(-0.2, 0.3, 7, dtype=torch.float64).expand(*batch, 7) + q = q.movedim(-1, 0).contiguous().movedim(0, -1).requires_grad_() + if batch == (2, 3): + assert not q.is_contiguous() + tcp = tcp.expand(*batch, 4, 4).clone() + tcp[..., 0, 3] += q.detach()[..., 0] + pose = model.forward_kinematics(q, tcp) + reference, _ = model.evaluate(q, tcp) + assert pose.shape == (*batch, 4, 4) + torch.testing.assert_close(pose, reference) + torch.testing.assert_close( + torch.autograd.grad(pose.sum(), q)[0], + torch.autograd.grad(reference.sum(), q)[0], + ) + + +@pytest.mark.parametrize("method", ["forward_kinematics", "evaluate"]) +def test_kinematics_rejects_mismatched_shapes_and_precision(method): + model, tcp = model_and_tcp() + evaluate = getattr(model, method) + q = torch.zeros(2, model.dof, dtype=torch.float64) + with pytest.raises(ValueError, match="shape"): + evaluate(q[..., :-1], tcp) + with pytest.raises(ValueError, match="shape"): + evaluate(q, tcp.expand(3, 4, 4)) + with pytest.raises(ValueError, match="dtype/device"): + evaluate(q.float(), tcp) + with pytest.raises(ValueError, match="dtype/device"): + evaluate(q, tcp.float()) + + def test_complete_ik_target_and_seed_gradcheck(): model, tcp = model_and_tcp() evaluate = lambda joints: model.evaluate(joints, tcp) diff --git a/tests/learning/layers/test_fk.py b/tests/learning/layers/test_fk.py new file mode 100644 index 000000000..752d24b0c --- /dev/null +++ b/tests/learning/layers/test_fk.py @@ -0,0 +1,96 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""FK-only network gradients and independent forward-call graphs.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +import torch + +from embodichain.compute.kinematics import SerialChainKinematics, pose_error +from embodichain.learning.layers import DifferentiableFKLayer + + +def make_fk(device: str = "cpu") -> DifferentiableFKLayer: + xml = ( + Path(__file__).resolve().parents[2] / "fixtures/kinematics/seven_dof.urdf" + ).read_text() + model = SerialChainKinematics.from_urdf( + xml, "link6", dtype=torch.float64, device=device + ) + tcp = torch.eye(4, dtype=torch.float64, device=device) + tcp[:3, 3] = tcp.new_tensor([0.1, -0.05, 0.08]) + return DifferentiableFKLayer(model, tcp).eval() + + +@pytest.mark.parametrize("device", ["cpu", pytest.param("cuda", marks=pytest.mark.gpu)]) +def test_fk_pose_loss_reaches_network_without_explicit_jacobian(monkeypatch, device): + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA unavailable") + fk = make_fk(device) + predictor = torch.nn.Linear(2, fk.dof, dtype=torch.float64, device=device) + with torch.no_grad(): + predictor.weight.copy_( + torch.linspace(-0.1, 0.2, 14, device=device).reshape(7, 2) + ) + predictor.bias.fill_(0.3) + # FK must evaluate joints even when they are outside configured limits. + fk.lower.fill_(-0.1) + fk.upper.fill_(0.1) + observation = torch.tensor( + [[-0.2, 0.3], [0.4, -0.1]], dtype=torch.float64, device=device + ) + with torch.no_grad(): + joints = predictor(observation) + reference, _ = fk.evaluate(joints) + target = fk(joints + 0.04) + + def unexpected_jacobian(*args, **kwargs): + raise AssertionError("FK-only calls must not construct a geometric Jacobian") + + monkeypatch.setattr(fk.model.chain, "jacobian_tensor", unexpected_jacobian) + network = torch.nn.Sequential(predictor, fk) + pose = network(observation) + torch.testing.assert_close(pose, reference) + position, rotation = pose_error(pose, target) + loss = (position.square().sum(-1) + rotation.square().sum(-1)).mean() + loss.backward() + for parameter in predictor.parameters(): + assert parameter.grad is not None + assert torch.isfinite(parameter.grad).all() and parameter.grad.norm() > 0 + assert not list(fk.parameters()) + + +def test_fk_repeated_calls_preserve_outputs_and_batch_gradients(): + fk = make_fk() + q = torch.linspace(-0.2, 0.3, 14, dtype=torch.float64).reshape(2, 7) + q.requires_grad_() + original = q.detach().clone() + first = fk(q) + snapshot = first.detach().clone() + second = fk(q + 0.02) + other_grad = torch.autograd.grad(first[0].sum(), q, retain_graph=True)[0] + torch.testing.assert_close(other_grad[1], torch.zeros_like(q[1])) + actual = torch.autograd.grad(first.sum() + second.sum(), q)[0] + expected = torch.autograd.grad( + fk.evaluate(q)[0].sum() + fk.evaluate(q + 0.02)[0].sum(), q + )[0] + torch.testing.assert_close(actual, expected) + torch.testing.assert_close(first, snapshot) + torch.testing.assert_close(q, original) diff --git a/tests/learning/layers/test_ik.py b/tests/learning/layers/test_ik.py index 84c5ce1f7..388ee9332 100644 --- a/tests/learning/layers/test_ik.py +++ b/tests/learning/layers/test_ik.py @@ -27,7 +27,7 @@ import torch from embodichain.compute.kinematics import DifferentiableIKCfg, SerialChainKinematics -from embodichain.learning.layers import DifferentiableIKLayer, KinematicsLayer +from embodichain.learning.layers import DifferentiableIKLayer, DifferentiableFKLayer def make_layer(dtype=torch.float64, device="cpu", iterations=4, enforce=True): @@ -38,7 +38,7 @@ def make_layer(dtype=torch.float64, device="cpu", iterations=4, enforce=True): tcp = torch.eye(4, device=device, dtype=dtype) tcp[:3, 3] = tcp.new_tensor([0.1, 0.02, -0.03]) return DifferentiableIKLayer( - KinematicsLayer(model, tcp), + DifferentiableFKLayer(model, tcp), DifferentiableIKCfg(max_iterations=iterations, enforce_joint_limits=enforce), ) @@ -51,7 +51,10 @@ def find_spec(self, fullname, *args): if any(fullname == n or fullname.startswith(n + '.') for n in ('embodichain.lab', 'embodichain.learning.rl', 'dexsim', 'warp', 'wandb', 'tensordict')): raise AssertionError(fullname) sys.meta_path.insert(0, Block()) -from embodichain.learning.layers import KinematicsLayer, DifferentiableIKLayer +from embodichain.learning.layers import DifferentiableFKLayer, DifferentiableIKLayer +from embodichain.learning.layers import KinematicsLayer +from embodichain.learning.layers.kinematics import KinematicsLayer as LegacyKinematicsLayer +assert KinematicsLayer is LegacyKinematicsLayer is DifferentiableFKLayer from embodichain.learning.models import MLP, EmpiricalNormalizer, build_mlp_from_cfg network = build_mlp_from_cfg({"type": "mlp", "network_cfg": {"hidden_sizes": [4]}}, 3, 2) assert isinstance(network, MLP) @@ -149,6 +152,15 @@ def test_checkpoint_dtype_and_eval_lifecycle(): target = layer.kinematics(q.detach() + 0.05) expected = layer(target, q) state = deepcopy(layer.state_dict()) + # These are the pre-rename checkpoint keys; the FK layer remains registered + # under "kinematics" and has no extra wrapper or extra buffer namespace. + assert set(state) == { + "_extra_state", + "kinematics.tcp", + "kinematics.lower", + "kinematics.upper", + "kinematics._extra_state", + } rebuilt = make_layer(dtype=torch.float64, iterations=1) rebuilt.load_state_dict(state) torch.testing.assert_close(expected, rebuilt(target, q)) diff --git a/tests/test_agent_context_map.py b/tests/test_agent_context_map.py index 105c569c0..c5aa65211 100644 --- a/tests/test_agent_context_map.py +++ b/tests/test_agent_context_map.py @@ -121,6 +121,8 @@ def test_representative_queries_route_against_the_repository_map() -> None: "Newton physics backend config 在哪里？": ["simulation-system"], "可微环境 APG 如何重置？": ["differentiable-env"], "可微 IK layer 的梯度": ["learning-layers"], + "可微 FK layer 的梯度": ["learning-layers"], + "DifferentiableFKLayer": ["learning-layers"], "DifferentiableIKLayer": ["learning-layers"], "sim.scene 的场景布局服务在哪里？": ["scene-services"], "程序化场景如何生成？": ["scene-services"], From 3e6b9a0e448619895d03e0b2dbfe4a52cad82b4f Mon Sep 17 00:00:00 2001 From: yuecideng Date: Thu, 29 Oct 2026 13:50:49 +0000 Subject: [PATCH 2/3] fix(learning): repair CI and remove legacy FK compatibility --- agent_context/MAP.yaml | 2 +- .../topics/learning-layers/learning-layers.md | 7 +- .../embodichain/embodichain.learning.rst | 7 +- .../overview/learning/differentiable_ik.md | 9 - .../learning/differentiable_kinematics.md | 14 +- embodichain/learning/layers/__init__.py | 3 +- embodichain/learning/layers/kinematics.py | 25 --- examples/learning/rigidformer/__main__.py | 3 +- examples/learning/rigidformer/_demo.py | 2 + .../kinematic_solver/run_benchmark.py | 176 +----------------- tests/learning/layers/test_ik.py | 3 - 11 files changed, 20 insertions(+), 231 deletions(-) delete mode 100644 docs/source/overview/learning/differentiable_ik.md delete mode 100644 embodichain/learning/layers/kinematics.py diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index c0705d8a0..a3b6aec66 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -292,7 +292,7 @@ topics: - id: learning-layers title: Differentiable Learning Layers aliases: [learning layers, differentiable FK layer, differentiable IK layer, 可微 FK, 可微 IK, 训练层, 学习模块分层] - keywords: [DifferentiableFKLayer, DifferentiableIKLayer, KinematicsLayer, SerialChainKinematics, DifferentiableIKCfg, + 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: diff --git a/agent_context/topics/learning-layers/learning-layers.md b/agent_context/topics/learning-layers/learning-layers.md index 710eeaf0b..5edb89c95 100644 --- a/agent_context/topics/learning-layers/learning-layers.md +++ b/agent_context/topics/learning-layers/learning-layers.md @@ -13,8 +13,7 @@ returns TCP poses and geometric Jacobians from one FK pass, retaining dJ/dq. `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. -`KinematicsLayer` remains a direct alias at both its original module and package -import paths. IK keeps its `kinematics` child name for checkpoint compatibility. +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 @@ -85,6 +84,4 @@ Tests live under `tests/compute/kinematics/` and `tests/learning/layers/`. Validate FK, dJ/dq, full IK and recurrent-state gradients separately; include nonzero TCP offsets, repeated solves, checkpoint migration and import boundaries. GPU validation runs serially with `--run-gpu`. The unified kinematic benchmark's `differentiable` mode measures training and -forward-only baselines, including FK/Jacobian and DLS component costs. Its explicit -`fk` mode compares pure FK with pose extraction through `evaluate`, including -matching pose-loss gradients and separately measured warmed memory. +forward-only baselines, including FK/Jacobian and DLS component costs. diff --git a/docs/source/api_reference/embodichain/embodichain.learning.rst b/docs/source/api_reference/embodichain/embodichain.learning.rst index 17f6790dd..549ba7249 100644 --- a/docs/source/api_reference/embodichain/embodichain.learning.rst +++ b/docs/source/api_reference/embodichain/embodichain.learning.rst @@ -38,8 +38,7 @@ 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`` composes an FK layer: its ``forward`` returns joints for network composition, while ``solve`` returns -residuals and a success mask. ``KinematicsLayer`` remains a direct compatibility -alias of ``DifferentiableFKLayer``, including its original module import path. +residuals and a success mask. Both retain input gradients in training and evaluation mode. Callers own rollout state, optimizers, and detach boundaries. @@ -49,14 +48,10 @@ state, optimizers, and detach boundaries. DifferentiableFKLayer DifferentiableIKLayer - KinematicsLayer .. automodule:: embodichain.learning.layers.fk :members: -.. automodule:: embodichain.learning.layers.kinematics - :members: - .. automodule:: embodichain.learning.layers.ik :members: diff --git a/docs/source/overview/learning/differentiable_ik.md b/docs/source/overview/learning/differentiable_ik.md deleted file mode 100644 index 346cfd2e3..000000000 --- a/docs/source/overview/learning/differentiable_ik.md +++ /dev/null @@ -1,9 +0,0 @@ ---- -orphan: true ---- - -# Differentiable FK and IK - -This guide has moved to {doc}`differentiable_kinematics`, which covers -shared setup, FK and IK training, gradient/state contracts, examples and -measurement. diff --git a/docs/source/overview/learning/differentiable_kinematics.md b/docs/source/overview/learning/differentiable_kinematics.md index 6a0f672d9..0ce24254d 100644 --- a/docs/source/overview/learning/differentiable_kinematics.md +++ b/docs/source/overview/learning/differentiable_kinematics.md @@ -158,9 +158,8 @@ check still uses physical limits. Numerical options belong to dict. Checkpoints store TCP/limit buffers, geometry identity and, for IK, numerical settings. `load_state_dict(assign=True)` also migrates fixed geometry to the loaded buffers' device and dtype, including nested loads. - `KinematicsLayer` remains a direct alias of `DifferentiableFKLayer` at its - original package and module import paths. IK retains the `kinematics` - constructor argument and child name for existing state dicts. + IK retains the `kinematics` constructor argument and child name for existing + state dicts. ## Examples and measurement @@ -177,16 +176,13 @@ The READMEs own runnable commands, dependencies, artifacts and evaluation protocols. These are kinematic examples; collision, contact, dynamics and physical execution are outside their validation scope. -The unified kinematic benchmark provides two relevant modes: +The unified kinematic benchmark measures differentiable IK: ```bash -python -m scripts.benchmark.robotics.kinematic_solver.run_benchmark -s fk python -m scripts.benchmark.robotics.kinematic_solver.run_benchmark -s differentiable ``` -`fk` compares pure FK with pose extraction through `evaluate` using identical -inputs and pose losses, checking pose/gradient agreement and separately measuring -warmed memory. `differentiable` measures IK forward and forward/backward passes, -FK/Jacobian and DLS components, and FEP/Pytorch forward-only baselines. Reports +`differentiable` measures IK forward and forward/backward passes, FK/Jacobian +and DLS components, and FEP/Pytorch forward-only baselines. Reports describe workloads, error metrics, timing and memory-accounting scope; benchmark results do not establish performance for a different task or device. diff --git a/embodichain/learning/layers/__init__.py b/embodichain/learning/layers/__init__.py index 96530a36c..6ab3b78c9 100644 --- a/embodichain/learning/layers/__init__.py +++ b/embodichain/learning/layers/__init__.py @@ -20,6 +20,5 @@ from .fk import DifferentiableFKLayer from .ik import DifferentiableIKLayer -from .kinematics import KinematicsLayer -__all__ = ["DifferentiableFKLayer", "DifferentiableIKLayer", "KinematicsLayer"] +__all__ = ["DifferentiableFKLayer", "DifferentiableIKLayer"] diff --git a/embodichain/learning/layers/kinematics.py b/embodichain/learning/layers/kinematics.py deleted file mode 100644 index 83dab9cbe..000000000 --- a/embodichain/learning/layers/kinematics.py +++ /dev/null @@ -1,25 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""Compatibility name for :class:`.fk.DifferentiableFKLayer`.""" - -from __future__ import annotations - -from .fk import DifferentiableFKLayer - -__all__ = ["KinematicsLayer"] - -KinematicsLayer = DifferentiableFKLayer diff --git a/examples/learning/rigidformer/__main__.py b/examples/learning/rigidformer/__main__.py index e44598c02..c5035e9a9 100644 --- a/examples/learning/rigidformer/__main__.py +++ b/examples/learning/rigidformer/__main__.py @@ -161,7 +161,6 @@ def evaluate(args: argparse.Namespace) -> None: from embodichain_models.dynamics.rigidformer import load_rigidformer_checkpoint from embodichain_models.dynamics.rigidformer.data import _verified_manifest from ._support._evidence import evidence - from ._support._evaluator import Evaluator if args.output.exists(): raise FileExistsError(args.output) @@ -169,6 +168,8 @@ def evaluate(args: argparse.Namespace) -> None: try: root = args.data_root.resolve() manifest = _verified_manifest(root) + from ._support._evaluator import Evaluator + model = load_rigidformer_checkpoint( args.checkpoint, device=args.device, diff --git a/examples/learning/rigidformer/_demo.py b/examples/learning/rigidformer/_demo.py index 64f4bce6d..9398cd64c 100644 --- a/examples/learning/rigidformer/_demo.py +++ b/examples/learning/rigidformer/_demo.py @@ -257,6 +257,7 @@ def run_demo(args: argparse.Namespace) -> None: ) from embodichain.lab.sim.sensors import CameraCfg from embodichain.lab.sim.shapes import SphereCfg + from embodichain.lab.visualization import VisualizationCfg from embodichain_models.dynamics.rigidformer import load_rigidformer_checkpoint from embodichain_models.dynamics.rigidformer.data import _verified_manifest @@ -300,6 +301,7 @@ def run_demo(args: argparse.Namespace) -> None: width=_WIDTH, height=_HEIGHT, render_cfg=RenderCfg(renderer="hybrid"), + visualization=VisualizationCfg(backend="none"), physics_cfg=NewtonPhysicsCfg( physics_dt=1 / _PHYSICS_HZ, num_substeps=2, diff --git a/scripts/benchmark/robotics/kinematic_solver/run_benchmark.py b/scripts/benchmark/robotics/kinematic_solver/run_benchmark.py index 89afbce64..ec46e286b 100644 --- a/scripts/benchmark/robotics/kinematic_solver/run_benchmark.py +++ b/scripts/benchmark/robotics/kinematic_solver/run_benchmark.py @@ -26,7 +26,6 @@ from __future__ import annotations import argparse -from collections.abc import Callable from importlib.metadata import version import os import time @@ -87,11 +86,11 @@ def _parse_args() -> argparse.Namespace: "--solvers", "-s", nargs="+", - choices=(*SUPPORTED_SOLVERS, "franka", "ssik", "fk", "all"), + choices=(*SUPPORTED_SOLVERS, "franka", "ssik", "all"), default=["all"], help=( "Solvers to benchmark. Use one or more of: opw, pytorch, ur, fep, differentiable, " - "franka, ssik, fk, all. Franka, SSIK and FK comparisons run only when selected explicitly. " + "franka, ssik, all. Franka and SSIK comparisons run only when selected explicitly. " "Default: all." ), ) @@ -222,13 +221,13 @@ def _normalize_selected_solvers(selected_solvers: list[str] | None) -> set[str]: """Normalize selected solver names to a canonical set.""" if not selected_solvers or "all" in selected_solvers: explicit_comparisons = set(selected_solvers or ()).intersection( - ("franka", "ssik", "fk") + ("franka", "ssik") ) return set(SUPPORTED_SOLVERS) | explicit_comparisons return { solver for solver in selected_solvers - if solver in (*SUPPORTED_SOLVERS, "franka", "ssik", "fk") + if solver in (*SUPPORTED_SOLVERS, "franka", "ssik") } @@ -1555,155 +1554,6 @@ def call_baseline(): return perf_rows, metric_rows -def benchmark_differentiable_fk() -> ( - tuple[list[dict[str, object]], list[dict[str, object]]] -): - """Compare pure FK with pose extraction through FK/Jacobian evaluation. - - Use identical seven-DOF Franka inputs, a translated/rotated TCP, and pose - losses for both implementations. Time device-resident steady-state calls; - measure warmed memory separately, retaining one output until after sampling. - """ - from embodichain.compute.kinematics import SerialChainKinematics, pose_error - from embodichain.learning.layers import DifferentiableFKLayer - - model = SerialChainKinematics.from_urdf( - Path(get_data_path("Franka/Panda/PandaWithHand.urdf")).read_text(), - "fr3_hand_tcp", - "base", - ) - tcp = torch.tensor( - [[0, -1, 0, 0.1], [1, 0, 0, -0.05], [0, 0, 1, 0.08], [0, 0, 0, 1]], - dtype=torch.float32, - ) - generator = torch.Generator().manual_seed(2026) - workloads = [] - for count in (1, 64, 1024): - q = model.chain.low + ( - 0.25 + 0.5 * torch.rand(count, model.dof, generator=generator) - ) * (model.chain.high - model.chain.low) - target = model.forward_kinematics(q + 0.04, tcp) - workloads.append((q, target)) - - def measure( - call: Callable[[], torch.Tensor], device: torch.device - ) -> dict[str, object]: - def synchronize() -> None: - if device.type == "cuda": - torch.cuda.synchronize(device) - - for _ in range(5): - output = call() - del output - times = [] - for _ in range(20): - synchronize() - start = time.perf_counter() - output = call() - synchronize() - times.append((time.perf_counter() - start) * 1000) - del output - - # RSS measures retained memory, not CPU transient peak. Torch's CUDA - # counters cover allocations by these Torch/PK paths, not total VRAM. - synchronize() - process = psutil.Process() - cpu_before = process.memory_info().rss - cuda = device.type == "cuda" - gpu_before = torch.cuda.memory_allocated(device) if cuda else 0 - if cuda: - torch.cuda.reset_peak_memory_stats(device) - output = call() - synchronize() - cpu_delta = (process.memory_info().rss - cpu_before) / 1024**2 - gpu_delta = ( - (torch.cuda.memory_allocated(device) - gpu_before) / 1024**2 - if cuda - else "N/A" - ) - peak = torch.cuda.max_memory_allocated(device) if cuda else 0 - del output - return { - "cost_time_ms": float(np.median(times)), - "p90_time_ms": float(np.percentile(times, 90)), - "cpu_delta_mb": cpu_delta, - "gpu_delta_mb": gpu_delta, - "peak_gpu_mb": peak / 1024**2 if cuda else "N/A", - "peak_gpu_delta_mb": (peak - gpu_before) / 1024**2 if cuda else "N/A", - } - - devices = [torch.device("cpu")] - if torch.cuda.is_available(): - devices.append(torch.device("cuda", torch.cuda.current_device())) - perf_rows: list[dict[str, object]] = [] - metric_rows: list[dict[str, object]] = [] - for device in devices: - fk = DifferentiableFKLayer(model, tcp).to(device) - for cpu_q, cpu_target in workloads: - q = cpu_q.to(device).detach().requires_grad_() - target = cpu_target.to(device) - - def loss(pose: torch.Tensor) -> torch.Tensor: - position, rotation = pose_error(pose, target) - return (position.square().sum(-1) + rotation.square().sum(-1)).mean() - - reference = fk.evaluate(q)[0] - reference_gradient = torch.autograd.grad(loss(reference), q)[0] - reference = reference.detach() - for name, operation in ( - ("fk_only", fk), - ("fk_via_jacobian", lambda joints: fk.evaluate(joints)[0]), - ): - pose = operation(q) - gradient = torch.autograd.grad(loss(pose), q)[0] - torch.testing.assert_close(pose, reference) - torch.testing.assert_close(gradient, reference_gradient) - with torch.no_grad(): - position, rotation = pose_error(pose, reference) - position_norm = position.norm(dim=-1) - rotation_norm = rotation.norm(dim=-1) - common = { - "sample_size": q.shape[0], - "impl": f"{name}_{device.type}", - "component": "fk_pose", - } - metric_rows.append( - { - **common, - "success_rate": ( - (position_norm < 1e-6) & (rotation_norm < 1e-6) - ) - .float() - .mean() - .item(), - "translation_err_mm": position_norm.mean().item() * 1000, - "rotation_err_deg": rotation_norm.mean().item() * 180 / np.pi, - "gradient_max_abs_error": (gradient - reference_gradient) - .abs() - .max() - .item(), - } - ) - del pose, gradient - for mode in ("no_grad", "forward", "forward_backward"): - - def call() -> torch.Tensor: - with torch.set_grad_enabled(mode != "no_grad"): - pose = operation(q) - if mode == "forward_backward": - return torch.autograd.grad(loss(pose), q)[0] - return pose - - measured = measure(call, device) - perf_rows.append({**common, "mode": mode, **measured}) - print( - f" {name} {device.type} n={q.shape[0]} {mode}: " - f"{measured['cost_time_ms']:.3f} ms", - flush=True, - ) - return perf_rows, metric_rows - - def run_all_benchmarks( selected_solvers: list[str] | None = None, *, @@ -1734,8 +1584,6 @@ def run_all_benchmarks( print( "- Franka comparison: default FEP search versus former Pytorch with 30 seeds." ) - if "fk" in solvers_to_run: - print("- FK comparison: pure FK versus pose extraction through FK/Jacobian.") perf_rows: list[dict[str, object]] = [] metric_rows: list[dict[str, object]] = [] @@ -1770,11 +1618,6 @@ def run_all_benchmarks( perf_rows.extend(diff_perf) metric_rows.extend(diff_metrics) - if "fk" in solvers_to_run: - fk_perf, fk_metrics = benchmark_differentiable_fk() - perf_rows.extend(fk_perf) - metric_rows.extend(fk_metrics) - ssik_notes: list[str] = [] if "ssik" in solvers_to_run: from scripts.benchmark.robotics.kinematic_solver.ssik_benchmark import ( @@ -1800,16 +1643,9 @@ def run_all_benchmarks( leaderboard_rows=leaderboard_rows, notes=[ *ssik_notes, - "Legacy CPU/GPU memory fields are deltas measured around timed calls; SSIK and FK use separate warmed memory passes. SSIK names Torch and Warp fields separately.", - f"Runtime: Torch {torch.__version__}, CUDA {torch.version.cuda}, pytorch_kinematics {version('pytorch_kinematics')}, CPU threads {torch.get_num_threads()}, GPU {torch.cuda.get_device_name() if torch.cuda.is_available() else 'unavailable'}. Numeric benchmark inputs use float32.", + "Legacy CPU/GPU memory fields are deltas measured around timed calls; SSIK uses a separate warmed memory pass with separately named Torch and Warp fields.", + f"Runtime: Torch {torch.__version__}, pytorch_kinematics {version('pytorch_kinematics')}, CPU threads {torch.get_num_threads()}, GPU {torch.cuda.get_device_name() if torch.cuda.is_available() else 'unavailable'}. Numeric benchmark inputs use float32.", "This report contains exactly three tables: Time & Memory, Success & Other Metrics, and Leaderboard.", - *( - [ - "FK mode: float32 Franka with translated/rotated TCP, batches 1/64/1024 and CPU RNG seed 2026. Inputs and targets are generated once on CPU, then moved before timing. Median and p90 completed-call wall latency use 20 synchronized repetitions after five warmups. no_grad measures inference; forward retains the graph; forward_backward rebuilds the graph and differentiates the same position/rotation loss without accumulating leaf gradients. Memory uses one separate warmed call retaining its output; CPU delta is retained RSS, not transient peak. CUDA peak_gpu_mb is total Torch peak allocated memory; peak_gpu_delta_mb is additional peak above the warmed baseline. CPU-only CUDA metrics are N/A. FK success_rate is pose agreement within 1e-6 m/rad of evaluate()[0]; gradient_max_abs_error compares input gradients. Geometry construction and transfers are excluded." - ] - if "fk" in solvers_to_run - else [] - ), *( [ "Differentiable mode uses shared CPU-generated Franka targets/seeds on CPU and CUDA. DLS: 4/8/16 steps, damping .05, step size .5, hard limits. FEP: redundancy search. Pytorch: two seeds, 64 iterations, lr .5. Timings are medians of three synchronized calls after two warmups, excluding data preparation and RNG reset to seed 2026; each backward call rebuilds its graph. Pose errors include all rows, including failures. Component rows are excluded from the success leaderboard. Peak GPU memory is absolute allocated Torch memory; Warp allocations in FEP are not captured by that counter." diff --git a/tests/learning/layers/test_ik.py b/tests/learning/layers/test_ik.py index 388ee9332..49b5531ea 100644 --- a/tests/learning/layers/test_ik.py +++ b/tests/learning/layers/test_ik.py @@ -52,9 +52,6 @@ def find_spec(self, fullname, *args): raise AssertionError(fullname) sys.meta_path.insert(0, Block()) from embodichain.learning.layers import DifferentiableFKLayer, DifferentiableIKLayer -from embodichain.learning.layers import KinematicsLayer -from embodichain.learning.layers.kinematics import KinematicsLayer as LegacyKinematicsLayer -assert KinematicsLayer is LegacyKinematicsLayer is DifferentiableFKLayer from embodichain.learning.models import MLP, EmpiricalNormalizer, build_mlp_from_cfg network = build_mlp_from_cfg({"type": "mlp", "network_cfg": {"hidden_sizes": [4]}}, 3, 2) assert isinstance(network, MLP) From ed496987b3eb2927ac00b9e95222691da7d17afe Mon Sep 17 00:00:00 2001 From: yuecideng Date: Thu, 29 Oct 2026 23:22:53 +0800 Subject: [PATCH 3/3] Fix optional dependencies in RigidFormer GPU recovery test Skip the GPU recovery case when roma or einops is unavailable, consistent with the neighboring optional RigidFormer tests. With dependencies installed, the existing exact EMA/recovery assertions still execute unchanged. Validation: 25 RigidFormer CPU/GPU tests and 59 FK/IK/import tests passed. Missing roma and missing einops were each verified to skip the case. Black 26.10.1, API docs coverage (3287/3287), and diff checks passed. Context review: no update needed; this test-only correction restores optional-dependency handling without changing documented model or training contracts. Existing behavioral coverage is reused; no redundant tests added. --- tests/models/dynamics/test_rigidformer_training.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/models/dynamics/test_rigidformer_training.py b/tests/models/dynamics/test_rigidformer_training.py index 32a388c98..2daebfaa4 100644 --- a/tests/models/dynamics/test_rigidformer_training.py +++ b/tests/models/dynamics/test_rigidformer_training.py @@ -421,6 +421,8 @@ def test_compact_windows_preserve_scene_stride_and_boundaries(random_sampling): @pytest.mark.gpu def test_recipe_resume_crosses_epoch_and_warmup_with_exact_ema(tmp_path): + pytest.importorskip("roma") + pytest.importorskip("einops") from embodichain_models.dynamics.rigidformer.sl_recipe import ( build_rigidformer_trainer, )