PaperlayerView exact revision on arXiv

Source-linked AI summary

Screw Attention: Rigid-Body Algebra Inside a Transformer

Aly Magassouba

arXiv:2610.00904v1cs.ROcs.AI

TL;DR

Learned manipulation policies must rediscover spatial relations that rigid-body mechanics supplies in closed form. Screw Attention builds those relations into a transformer through frame-transported messages, and it performs strongly against learned controls while remaining robust to per-link frame changes.

  • Problem

    Learned policies must recover spatial relations supplied directly by rigid-body mechanics, while analytic methods cannot readily represent contact, friction, deformation, uncertain payloads, or non-fixed set-point objectives.

    Source: §I · p. 1

  • Method

    Screw Attention uses body tokens, pairwise physical relations, and frame-transported messages, with a gated learned residual optionally added to an analytic controller.

    Source: §I · p. 2, C. Readouts and the residual · p. 6

  • Results

    Screw Attention outperforms same-size graph, transformer, and flat networks on LIBERO-Spatial and is the only model unchanged under per-link frame changes, while controls fall to 0.0–2.3%.

    Source: §I · p. 2

  • Takeaways & Limitations

    Across the studied tasks, geometry is decisive when a task requires relations between coordinate frames that an analytic controller does not otherwise supply.

    Source: §I · p. 2

    Validation was only in simulation; translation noise was tested, but rotation errors were not.

    Source: §VI · p. 12

Abstract

from arXiv · show

Learned manipulation policies rediscover from data the spatial relations that rigid-body mechanics supplies in closed form. This costs data, and it leaves the policies fragile to geometric changes in the scene. We present Screw Attention, a transformer layer in which the relation between two bodies is a spatial transform rather than a graph edge. Every token is a body with a pose. Each pair of tokens carries the relative pose and, for robot joints, the joint screw. Messages are transported along this relation into the receiver's frame, while the attention scores see only frame-invariant quantities. By construction, the messages are equivariant to an independent change of frame at every token, and a single layer can express the velocity recursion of rigid-body mechanics. On simulated manipulation tasks, Screw Attention matches or outperforms controls of the same size, including graph, transformer and flat networks on LIBERO-Spatial. With 16,162 parameters it reaches 97.3% on LIBERO-Spatial from object poses (without images or language), above a flat network with 27x more parameters. Under a change of per-link frame convention its success is unchanged, while every other learned network falls below 3%. Placed on an analytic controller as a gated residual, it raises insertion success by 17.3 points. It is unaffected by pose noise up to 10,mm and by joint offsets within the factory calibration of a Franka arm. These results suggest a criterion: geometry is decisive when the task requires relations between frames that no other part of the system supplies. Code and trained policies will be released.

I. INTRODUCTION

Learned manipulation policies must recover spatial relations that rigid-body mechanics specifies, while existing approaches preserve only parts of the available structure. Screw Attention represents body-to-body relations physically and transports messages between frames, enabling independent per-link frame equivariance and exact expression of velocity recursion.

I. INTRODUCTION

  • Learned policies can fail when object placements or robot configurations change, and policies tied to arbitrary link-frame conventions can fail on physically identical robots described differently.Rigid-body mechanics supplies kinematic relations in closed form, while analytic methods cannot easily express contact, friction, deformation, uncertain payloads, or non-fixed set-point objectives.Source: §I · pp. 1–2

II. RELATED WORK

  • Earlier policy transformers encode robot connectivity, equivariant policies structure scene geometry, and kinematics-structured policies carry link poses or transforms but process them with learned maps.These approaches leave physical relations to inference or learned reinterpretation rather than directly encoding their spatial-algebra content.Source: §I · pp. 1–2, §II · pp. 2–3

I. INTRODUCTION

  • The layer represents each body as a token and relates token pairs through physical relations, transporting messages into the receiver’s frame rather than relying on abstract graph edges.The relations are available from the robot model or perception, and the layer can express motion propagation along the robot rather than learning it from data.Source: Figure 1 · p. 1, §I · p. 2
  • Screw Attention is equivariant to an independent change of frame on every link, unlike existing equivariant architectures that handle only a global frame change.The global change is a special case of the layer’s symmetry, and the construction applies the same geometric typing to robot bodies and scene objects.Source: §I · p. 2, §II · p. 3

II. RELATED WORK

  • Screw Attention keeps transforms, screws, and inertias as model quantities in frame-commuting operations, making velocity recursion an exact special case rather than an approximation.XMoP compensates for frame dependence with frame-randomising augmentation and reports losing 68.1% of its success rate on unseen embodiments without it.Source: §II · p. 3
  • Physics-based objectives, separate models, and residual policies constrain other parts of the system and are compatible with constraining the network’s function class.An analytic decoder can handle redefinitions of end-effector and camera frames when its image and language inputs do not change; per-link frame changes alter the quantities read by this policy.Source: §II · p. 3

III. METHOD

Screw Attention represents each robot link and scene object as a token with scalar and geometric channels. Pairwise spatial transforms and joint screws provide the relations used to transport geometric messages and produce commands.

III. METHOD

  • Each robot body and scene object is a token, updated in scalar and geometric lanes that share attention weights, before readout produces a command.Geometric messages are transported into the receiver’s frame, and an optional analytic-controller output can be added at readout.Source: §III · p. 3, Figure 2 · p. 4

A. Tokens and the pair term

  • A twist combines angular and linear velocity, and a wrench combines moment and force; both are six-dimensional screws expressed in a body frame.The spatial transform Xij maps a twist from body j’s frame into body i’s frame, while a wrench uses the dual transform.Source: A. Tokens and the pair term · p. 3
  • Token inputs include frame-independent features and twists expressed in the token’s frame, which are embedded into scalar features and geometric channels.Examples include joint position and velocity, mass, joint twist, gravity, and goal relations; geometric channels are weighted sums of input twists.Source: A. Tokens and the pair term · pp. 3–4
  • For force readout, tokens can also carry wrench channels and spatial inertia; object inertia is set to zero.Robot-body inertia is supplied by the robot model, and wrench channels are zero when no wrench is measured.Source: A. Tokens and the pair term · p. 4
  • Each ordered token pair carries Xij for the relative pose, and robot pairs also carry the sender joint’s motion subspace Sj.Robot transforms are computed from joint positions by forward kinematics; object transforms come from perceived poses, and pair terms are not learned.Source: A. Tokens and the pair term · pp. 3–4

B. The attention layer

Screw Attention uses frame-invariant pairings and kinematic masks to score interactions, then transports geometric messages into each receiving body’s frame. Its gates and channel updates can express rigid-body computations while preserving the geometry of the messages.

Pairings

  • The Killing and Klein pairings provide frame-invariant comparisons between twists, including rotation and sliding rates relative to a joint axis.The standard dot product on twists is frame-dependent, while κ and λ are preserved under frame changes.Source: B. The attention layer · p. 4

Masks

  • Outward heads attend to a body and its ancestors, while inward heads attend to it and its descendants, mirroring the two Newton–Euler sweeps.Robot tokens also attend to scene tokens; scene tokens attend only to themselves.Source: B. The attention layer · p. 5

Attention weights

  • Attention scores combine learned query–key compatibility with given geometric and kinematic features and pairings of transported channels.The pair features describe relative frame placement, joint-axis relations, and kinematic-tree structure.Source: B. The attention layer · pp. 4–5

Updates

  • Geometric messages mix channels, transform them into the receiver’s frame, and use a scalar-feature gate to scale each channel.Learned matrices act on channel indices, while the given spatial transform acts on the six geometric components.Source: B. The attention layer · p. 5

The gate

  • A gated outward-head update can reproduce the velocity recursion by compensating for the softmax’s averaging of ancestor messages.The gate scales by the number of ancestors, allowing the update to match the unnormalized sum.Source: B. The attention layer · p. 5

Wrench channels

  • Wrench channels add inertia-induced wrenches before exchanging messages through dual transforms, following the Newton–Euler rigid-body equation.Without inertia coupling, initially zero wrench channels would remain zero under the linear message update.Source: B. The attention layer · pp. 5–6

C. Readouts and the residual

Screw Attention offers motion, force, and twist readouts suited to different commands, and can also add a gated learned residual to an analytic controller. Shared readout weights support arms with different numbers of joints.

Torque readouts

  • Motion and force readouts produce joint torques by pairing each joint screw with geometric or wrench channels, respectively.The motion readout contains the Jacobian-transpose law; the force readout is the default for torque commands.Source: C. Readouts and the residual · p. 6

Shared readout

  • The same readout weights apply across joints, so the network can be used unchanged on arms with different numbers of degrees of freedom.No readout parameter is tied to a particular joint.Source: C. Readouts and the residual · p. 6

Twist readout

  • The twist readout emits a bounded end-effector twist in that token’s frame, which is transformed outside the network for the controller.Its angular and linear components are each bounded by their norm.Source: C. Readouts and the residual · p. 6

Learning alone or with a controller

  • A gated residual combines the analytic controller’s command with a learned command from the layer, and closing the gate recovers the controller exactly.The model term has weight one; the residual can learn what the model omits, such as contact and friction.Source: C. Readouts and the residual · p. 6

D. Representability

Screw Attention can represent classical rigid-body computations with hand-set weights rather than requiring a policy to learn them from data. A three-layer stack represents inverse dynamics, though the paper’s policies use two layers.

Representability

  • One outward head reproduces the velocity recursion with three nonzero weights and a gate, carrying ancestor joint twists to the target link.The construction uses uniform attention over ancestors and a gate that removes softmax normalization.Source: D. Representability · p. 6, Figure 3 · p. 7
  • One inward head and one readout weight reproduce the Jacobian-transpose law by carrying the end-effector orientation error to ancestor joints.The figure illustrates this construction for the angular Jacobian; the same construction with wrench channels gives τ = J⊤F.Source: D. Representability · pp. 6–7, Figure 3 · p. 7
  • A three-layer stack reproduces inverse dynamics, with one layer for each dependent velocity, acceleration, and wrench stage.The paper’s policies use two layers, and full inverse dynamics is left to future work.Source: D. Representability · p. 7

E. Frame invariance

Independent changes to token frames preserve scalar features and pairings, while the spatial operations transform covariantly. The layer’s outputs therefore transform consistently with its inputs, with a qualification for the twist readout.

Frame invariance

  • Scalar features, the κ and λ pairings, and twist–wrench power remain unchanged when each token’s frame changes independently.Learned matrices act only on channel indices, while transport, inertia coupling, and cross products are covariant.Source: E. Frame invariance · p. 7
  • The layer outputs transform exactly as their inputs do under independent token-frame changes, and the properties are verified numerically.Source: E. Frame invariance · p. 7
  • Torque commands are invariant to frame changes, but the twist readout is invariant only to end-effector-frame rotations because its norm bound depends on the expression point.Source: E. Frame invariance · p. 7

IV. EXPERIMENTS

The experiments evaluate Screw Attention on three simulated robot tasks and the LIBERO-Spatial benchmark. The tasks cover orientation control, selecting a reachable approach frame, contact insertion, and object placement.

A. Tasks

  • The evaluation spans reorientation, affordance selection, peg insertion, and ten LIBERO-Spatial object-placement tasks.LIBERO-Spatial placement success is evaluated on 500 canonical episodes; the policy uses simulator poses rather than camera or language inputs.Source: A. Tasks · pp. 7–8
  • Reorientation tests whether the arm can bring a payload of randomized mass to a commanded orientation, with goals drawn from one hemisphere of SO(3) in training and both hemispheres at evaluation.Success requires angle error below 0.1 rad for 20 steps over 512 fixed episodes.Source: A. Tasks · p. 7, Figure 4 · p. 8
  • Affordance selection requires choosing and reaching one of eight candidate approach frames without collision, although only an unknown subset is reachable from the current configuration.The policy is not told which candidates are feasible; evaluation uses 1,024 fixed scenes.Source: A. Tasks · p. 7, Figure 4 · p. 8
  • Peg insertion evaluates a learned torque residual when inserting a flange-welded peg into a fixed socket, including episodes inside and beyond the analytic controller’s reach.The metric is insertion success reported as a margin over the analytic controller on the same episodes.Source: A. Tasks · p. 7, Figure 4 · p. 8

B. Controls

The controls isolate geometric channels, transform assignments, and attention masks, and also compare against alternative network families. Training and configuration are specified separately for the robot tasks and LIBERO-Spatial.

B. Controls

  • The controls include scalar-only and shuffled-transform ablations, unconstrained geometric mixing, alternative attention masks, and flat, transformer, and graph networks.Task-specific controls additionally remove the scene screw for affordance selection or shuffle inertias for insertion.Source: B. Controls · p. 8
  • The pre-registered test on every robot task is whether Screw Attention outperforms both shuffled transforms and scalar only.The analytic controller and flat MLP are additional controls; the matched-budget transformer and graph network are used on LIBERO-Spatial.Source: B. Controls · p. 8

C. Training and configuration

  • Reorientation and affordance use behavior cloning followed by PPO, while insertion uses PPO alone; LIBERO-Spatial networks fit the benchmark demonstrations and are re-executed with object offsets up to 5 cm.Networks use two layers, four heads, and d = 32; task-specific readouts, channel counts, and parameter counts differ.Source: C. Training and configuration · p. 8

D. Does the relation pay?

Screw Attention leads the robot-task comparisons, while ablations show that geometry matters most when the rest of the control loop does not supply it. Frame-change tests also show that its equivariance preserves behavior across token-frame changes.

D. Does the relation pay?

  • 80.2±1.3% affordance success beats scalar only by 64.2 points and shuffled transforms by 73.9, while reorientation beats both controls by 96.5 points each.On affordance, Screw Attention exceeds the analytic controller’s 71.5% on every seed; on reorientation it is within 0.8 points of that controller.Source: D. Does the relation pay? · p. 9
  • Screw Attention raises insertion success 17.3 points above its analytic model term, with the largest gain, 25.9 points, on episodes beyond the controller’s reach.Its margin over shuffled transforms is unresolved: +2.7 points on three of four seeds.Source: D. Does the relation pay? · p. 9

E. Where does the geometry pay?

  • 8.1±4.2% affordance success after removing the scene screw and 6.3±2.2% after shuffling transforms show the task depends on both the candidate direction and its correct transport.Flat MLPs reach 0.8 ± 0.4% at matched budget and 4.7±1.8% at 405k parameters; scalar only reaches 16.0±3.6%.Source: E. Where does the geometry pay? · p. 9, Figure 5 · p. 9
  • Insertion loses only 6.5 points without geometric channels, while the 2.7-point shuffled-transform gap is unresolved and shuffled inertias cost 1.3 points.The trained policies changed their actions by under 0.2% when inertias were permuted, so they did not use inertia in this task.Source: E. Where does the geometry pay? · p. 9, Figure 5 · p. 9

F. What does the equivariance constraint buy?

  • A 0.6-point reorientation loss under changed link frames contrasts with a 51.7-point loss for the unconstrained network; other equivariant variants lose 0.4–6.5 points.The construction’s remaining frame-change effects pass through three pair features describing frame placement.Source: F. What does the equivariance… · pp. 9–10, Table II · p. 10
  • 97.7% LIBERO-Spatial success after changing every token’s frame, versus 97.3% canonically, contrasts with controls retaining only 0.0–2.7% of canonical success.The frame-change test isolates the transported path by rotating each token’s frame while keeping pair features canonical.Source: F. What does the equivariance… · p. 10, Figure 6 · p. 10
  • The equivariance constraint does not reduce demonstrations needed when all demonstrations share one frame convention, while the attention mask reduces sample needs and affects full-budget success.At full budget, adjacency and same-density random masks reach 75.5% and 51.9%, respectively, versus 98.0% without a mask.Source: F. What does the equivariance… · p. 10

G. LIBERO-Spatial

On LIBERO-Spatial, Screw Attention outperforms matched-budget networks and remains strong under several geometric perturbations. Its success declines with larger pose noise or joint offsets, including offsets beyond the cited factory-calibration range.

G. LIBERO-Spatial

  • 97.3% canonical LIBERO-Spatial success exceeds the matched-budget flat MLP’s 19.9%, transformer’s 75.8%, and graph network’s 84.7%.A flat MLP with 27.1 times more parameters reaches 95.2%; Screw Attention takes 0.66 ms per CPU forward pass versus 0.05 ms for the flat MLP.Source: G. LIBERO-Spatial · p. 10
  • Object displacements up to 5 cm lower success by 6.1–10.4 points for all networks except the matched flat MLP, whose canonical success is already low.Because the training data also vary object placement, this tests a harder distribution rather than robustness.Source: G. LIBERO-Spatial · pp. 10–11
  • 96.5% success at 10 mm pose noise is close to Screw Attention’s 97.3% canonical success, while success falls off by 40 mm.At 20 mm, Screw Attention loses 2.6 to 4.2 points; the analytic controller given the same poses loses 50.8 points.Source: G. LIBERO-Spatial · p. 11, Figure 8 · p. 11

H. Joint offsets

  • Within the cited factory-calibration range, Screw Attention loses 0.4 points at 0.25° and 1.7 points at 0.5° joint offsets; at 2° it loses 42.5 points.The offsets displace the hand by 5 and 10 mm on average at 0.25° and 0.5°, respectively; both Screw Attention and the analytic controller fall off sharply at 2°.Source: H. Joint offsets · p. 11

V. DISCUSSION

Screw Attention’s geometry helps most when a task depends on frame relations not supplied elsewhere, while its per-link frame equivariance improves robustness to changes in robot description. The study also identifies simulation, robot-model, and joint-structure boundaries for the approach.

A. Where the geometry contributes, and where it does not

  • 73.9 points separate correct from shuffled relations on affordance, where reachability must be inferred; insertion benefits less because its analytic controller already supplies the task geometry.The insertion residual nevertheless adds 17.3 points to the analytic controller.Source: A. Where the geometry contributes, and… · p. 11, §VII · p. 12

B. What the equivariance constraint provides, and what it does not

  • Under per-link frame changes, constrained networks retain nearly all success, while unconstrained networks lose about half on average and LIBERO controls lose nearly all.The attention mask, not equivariance, reduces the number of demonstrations needed; pose noise and object displacement affect learned networks similarly.Source: B. What the equivariance constraint… · pp. 11–12

C. Why the learned policies tolerate observation errors

  • Under pose noise and joint-angle offsets, learned policies lose less success than the analytic controller receiving the same corrupted observations.The analytic controller relies on precise discrete decisions, whereas learned policies map observations continuously; tolerance to pose noise is shared by learned networks and disappears beyond grasp tolerance.Source: C. Why the learned policies tolerate… · p. 12

VI. LIMITATIONS

  • Validation was limited to simulation using simulator object poses; deployment requires a pose estimator, and rotation-estimation errors were not tested.Photometric factors such as camera placement, lighting, background, and texture were outside the paper’s scope.Source: §VI · p. 12
  • The method depends on robot-model accuracy, but model errors were tested only as joint offsets, not as link-length or inertia errors.Source: §VI · p. 12
  • The construction assumes fixed-base serial-arm bodies driven by one revolute or prismatic joint, excluding multi-joint bodies and floating bases; inference is slowest and scales quadratically with token count.At 22 tokens, inference remains far below the control period.Source: §VI · p. 12

VII. CONCLUSION

  • On four tasks, Screw Attention matched or outperformed same-size learned networks; on LIBERO-Spatial it reached 97.3% with 16,162 parameters, above a flat network with 27 times more parameters.Its gated residual added 17.3 points to insertion success, and its frame change cost 0.6 points on reorientation versus 51.7 points for an unconstrained policy.Source: §VII · p. 12

APPENDIX A TRAINING DETAILS

The appendix records the simulation and training setup, checks that the spatial-algebra constructions reproduce rigid-body computations, and details how frame-dependent pair features enter attention scores. It also contrasts Screw Attention’s per-token frame symmetry with the global-frame symmetry of geometric-algebra transformers.

APPENDIX A TRAINING DETAILS

  • All networks were implemented in JAX and tasks simulated in MuJoCo; LIBERO policies used 979 demonstration episodes and 119,968 states with object offsets up to 5 cm.Insertion omitted behavior cloning because the model term already implements the demonstrator and the residual target is zero.Source: Appendix A · p. 14

APPENDIX B VERIFICATION OF THE CONSTRUCTION

  • Spatial-algebra quantities matched an independent rigid-body engine to 10^-13 or better, and one layer reproduced velocity recursion to 4×10^-15 on five kinematic trees.Hand-set weights also reproduced the Jacobian-transpose torque law to 10^-9 across chain, branched, and welded-body trees.Source: Appendix B · p. 14
  • With frame-dependent pair features removed, the layer commuted with independent per-body frame changes to 10^-10; mixing spatial components with Wg failed the tests.No learned parameter acts on the six spatial components.Source: Appendix B · p. 14

APPENDIX C RELATION TO GEOMETRIC-ALGEBRA TRANSFORMERS

  • GATr handles a common frame change, whereas Screw Attention’s per-token frame symmetry includes independent link-frame changes; the former cannot reproduce the latter’s changed relative transforms.Screw Attention trades algebraic generality for this per-token gauge, with the stated invariance limited to torque readouts and subject to frame-dependent pair features.Source: Appendix C · p. 15

APPENDIX D TERMS OF THE ATTENTION SCORE

  • The attention score uses fourteen pair features; only the first three depend on frame placement, and attention weights depend on frames only through those entries.For pairs involving an object, joint entries are zero and graph distance takes its maximum value, except for the object paired with itself.Source: Appendix D · p. 15

APPENDIX E CONTROLS AND THE MECHANISMS THEY INSTANTIATE

The appendix explains why controls reimplement published mechanisms at matched budgets rather than compare published results directly, and clarifies the pose-based evaluation scope. Across robustness and task analyses, it distinguishes frame-convention effects from settings where scene relations add little.

APPENDIX E CONTROLS AND THE MECHANISMS THEY INSTANTIATE

  • Controls reimplement published mechanisms at matched parameter budgets and shared training settings, not full published methods; the LIBERO graph control even receives true relative transforms.The graph control is better informed than ABD-NET, which learns its edge features and inertias.Source: Appendix E · pp. 15–16

APPENDIX F ROBUSTNESS ON LIBERO-SPATIAL

  • Screw Attention loses 6.8 points under a perturbed initial robot configuration, versus 38.6 for MINERVA and 44.9 for UniVLA; only drops are comparable across these pose- and image-based evaluations.The perturbation moves the robot while leaving objects at training positions.Source: Table VIII · p. 16, Appendix F · p. 16

APPENDIX E CONTROLS AND THE MECHANISMS THEY INSTANTIATE

  • The study uses simulator object poses rather than images, so its results are not directly comparable with image-based policies; deployment requires a pose estimator not counted in the reported parameters.Source: Appendix E · p. 16

APPENDIX G TRAINING FIT ON LIBERO-SPATIAL

  • Three of four controls fit LIBERO-Spatial training data better than Screw Attention, yet Screw Attention has the highest canonical success; controls retain at most 2.7% success after a frame change.The authors therefore evaluate results by rollout success rather than held-out fit.Source: Table IX · p. 16, Appendix G · p. 16

APPENDIX H WHERE THE SCENE RELATION DOES NOT HELP

  • With a reactive demonstrator, full-scene input differed from nearest-obstacle input by 2.3 ± 1.5 points, while single-pose task gains mainly came from encoding goal orientation as an end-effector screw.The single-pose score advantage was 0.318 over scalar only and 0.319 over shuffled transforms; object geometric channels were zero.Source: Appendix H · p. 17

Related papers

Loading 2610.00904v1…