diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index b81e03e7f..a3b6aec66 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, 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..5edb89c95 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,11 @@ DLS, SE(3) residuals, copied PK geometry, and shared configuration/results. These modules do not import learning or simulation. `SerialChainKinematics` returns TCP poses and geometric Jacobians from one FK pass, retaining dJ/dq. -`learning/layers/` owns `KinematicsLayer` and `DifferentiableIKLayer` Torch -module interfaces. `learning/models/` owns general MLPs, feature normalization +`learning/layers/fk.py` owns `DifferentiableFKLayer`: joints to root-frame TCP +poses through a pure FK path; `evaluate` also returns geometric Jacobians from +one FK pass. `DifferentiableIKLayer` composes this FK layer with fixed-step DLS. +IK keeps its `kinematics` child name for checkpoint compatibility. +`learning/models/` owns general MLPs, feature normalization and configuration-driven MLP construction. These shared components do not import RL, TensorDict or simulation. Concrete NeRD models, codecs and the shared-SL recipe live in the bundled `embodichain_models.dynamics.nerd` extension; @@ -20,8 +26,8 @@ policy registration/builders belong to `learning/rl/policies/`; RL code consumes the shared networks and normalizer. See [RL learning](../rl-learning/rl-learning.md) for those contracts. -Training, task-parameter optimization and the differentiable benchmark call -`DifferentiableIKLayer` directly. Use `torch.no_grad()` for forward-only inference. +Training, task-parameter optimization and the differentiable benchmark compose +FK or IK layers directly. Use `torch.no_grad()` for forward-only inference. There is no registered Robot solver adapter: callers provide root-frame targets and explicitly configure the layer's TCP and limits. Existing Robot solver binding and arena conversion remain in [IK solvers](../ik-solvers/ik-solvers.md). @@ -33,7 +39,9 @@ Supervised optimizer and offline-run ownership belongs to Targets are root-frame TCP matrices; seed batch axes match the target or contain only a joint axis. Training inputs must match model device and dtype. Both float32 -and float64 are supported. FK and Jacobian derivatives remain connected through +and float64 are supported. FK inputs follow the model's joint order; FK does not +project limits. Caller-owned pose losses can train joint predictors through FK. +FK and Jacobian derivatives remain connected through all iterations. Final tolerances produce boolean metadata; no convergence-based Python exit occurs. Failure returns the final iterate and continuous residuals. Hard limits have piecewise gradients; disabling projection does not disable the @@ -55,6 +63,13 @@ Model and layer callers own task losses and optimization. Formal task deployment belong to the task extension. See [RL learning](../rl-learning/rl-learning.md) for trainer/collector lifecycle. +`examples/learning/differentiable_fk/learn_pose_reaching.py` trains an example-local +joint predictor through fixed FK with TCP pose supervision and the shared SL +trainer. Generating joint solutions are discarded; validation selects weights +before independent local test evaluation. Its viewer compares initial and frozen +trained predictions using one CPU worker and rejects stale target edits. Example +smoke commands own launch/artifact checks; production tests own FK numerics. + `examples/learning/differentiable_ik/learn_grasp_yaw.py` trains an example-local conditional yaw predictor through fixed-step IK, preserving grasp position and downward approach. Validation alone selects checkpoints; independent target 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..549ba7249 100644 --- a/docs/source/api_reference/embodichain/embodichain.learning.rst +++ b/docs/source/api_reference/embodichain/embodichain.learning.rst @@ -32,10 +32,13 @@ Differentiable layers --------------------- ``learning.layers`` composes simulation-independent computation as Torch modules. -``KinematicsLayer`` owns fixed TCP and limit buffers; its copied PK geometry follows +``DifferentiableFKLayer.forward`` maps joints to root-frame TCP transforms without +constructing a geometric Jacobian; ``evaluate`` returns both in one FK pass. +The layer owns fixed TCP and limit buffers, and its copied PK geometry follows ``to(device/dtype)``. Reconstruct the same robot model before loading a checkpoint; -its geometry fingerprint is checked. ``DifferentiableIKLayer.forward`` returns joints -for network composition, while ``solve`` returns residuals and a success mask. +its geometry fingerprint is checked. ``DifferentiableIKLayer`` composes an FK layer: +its ``forward`` returns joints for network composition, while ``solve`` returns +residuals and a success mask. Both retain input gradients in training and evaluation mode. Callers own rollout state, optimizers, and detach boundaries. @@ -43,10 +46,10 @@ state, optimizers, and detach boundaries. .. autosummary:: - KinematicsLayer + DifferentiableFKLayer DifferentiableIKLayer -.. automodule:: embodichain.learning.layers.kinematics +.. automodule:: embodichain.learning.layers.fk :members: .. automodule:: embodichain.learning.layers.ik diff --git a/docs/source/overview/learning/differentiable_ik.md b/docs/source/overview/learning/differentiable_ik.md deleted file mode 100644 index 9f5cb9ec6..000000000 --- a/docs/source/overview/learning/differentiable_ik.md +++ /dev/null @@ -1,113 +0,0 @@ -# Differentiable IK - -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. - -## 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. diff --git a/docs/source/overview/learning/differentiable_kinematics.md b/docs/source/overview/learning/differentiable_kinematics.md new file mode 100644 index 000000000..0ce24254d --- /dev/null +++ b/docs/source/overview/learning/differentiable_kinematics.md @@ -0,0 +1,188 @@ +# 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. + 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 measures differentiable IK: + +```bash +python -m scripts.benchmark.robotics.kinematic_solver.run_benchmark -s differentiable +``` + +`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..6ab3b78c9 100644 --- a/embodichain/learning/layers/__init__.py +++ b/embodichain/learning/layers/__init__.py @@ -18,7 +18,7 @@ from __future__ import annotations -from .kinematics import KinematicsLayer +from .fk import DifferentiableFKLayer from .ik import DifferentiableIKLayer -__all__ = ["KinematicsLayer", "DifferentiableIKLayer"] +__all__ = ["DifferentiableFKLayer", "DifferentiableIKLayer"] diff --git a/embodichain/learning/layers/kinematics.py b/embodichain/learning/layers/fk.py similarity index 77% rename from embodichain/learning/layers/kinematics.py rename to embodichain/learning/layers/fk.py index 4701f5eea..ca07cd269 100644 --- a/embodichain/learning/layers/kinematics.py +++ b/embodichain/learning/layers/fk.py @@ -14,7 +14,7 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Torch module interface to fixed robot geometry, independent of simulation.""" +"""A differentiable forward-kinematics layer with fixed robot geometry.""" from __future__ import annotations @@ -25,10 +25,10 @@ from embodichain.compute.kinematics.serial_chain import SerialChainKinematics -__all__ = ["KinematicsLayer"] +__all__ = ["DifferentiableFKLayer"] -class KinematicsLayer(nn.Module): +class DifferentiableFKLayer(nn.Module): """Expose FK and differentiable Jacobians as a reusable network layer. Args: @@ -43,6 +43,8 @@ class KinematicsLayer(nn.Module): ``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__( @@ -101,11 +103,11 @@ def dof(self) -> int: """Number of movable joints in solver order.""" return self.model.dof - def _apply(self, fn: Any, recurse: bool = True) -> KinematicsLayer: + 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("KinematicsLayer supports float32 and 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 @@ -116,7 +118,7 @@ def _sync_model_after_load( # 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") + raise ValueError("DifferentiableFKLayer supports float32 and float64") if ( self.model.chain.device != self.tcp.device or self.model.chain.dtype != self.tcp.dtype @@ -135,15 +137,40 @@ def set_extra_state(self, state: dict[str, str]) -> None: ) def forward(self, qpos: torch.Tensor) -> torch.Tensor: - """Return root-frame TCP transforms for ``(..., dof)`` joints.""" - return self.evaluate(qpos)[0] + """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]: - """Return TCP transforms and root-frame geometric Jacobians.""" + """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, shape ``(..., 6, dof)``.""" + """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: 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/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/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 a016e865d..ec46e286b 100644 --- a/scripts/benchmark/robotics/kinematic_solver/run_benchmark.py +++ b/scripts/benchmark/robotics/kinematic_solver/run_benchmark.py @@ -1327,7 +1327,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 +1395,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", 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..49b5531ea 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,7 @@ 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.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 +149,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/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, ) 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"],