Pure Nx implementation of the FABRIK algorithm.
Provides both position-only solving (fabrik/5) and full frame-based
solving with orientation (fabrik_with_orientation/7).
Summary
Functions
FABRIK backward reaching pass.
Orientation-aware FABRIK backward reaching pass.
Run the FABRIK algorithm on a chain of points.
Run FABRIK with full orientation support.
FABRIK forward reaching pass.
Orientation-aware FABRIK forward reaching pass.
Extract positions tensor from frames.
Place a point at desired_distance from anchor, along the direction from
anchor toward point_to_move. The per-joint reaching step shared by both
passes; a defn so it composes into them and is reusable on its own.
Convert points tensor to frames with identity orientations.
Run the FABRIK iteration loop until convergence or max_iterations.
Run the orientation-aware FABRIK loop until convergence or max_iterations.
Stretch the chain straight toward an unreachable target.
Types
@type frames() :: %{positions: Nx.Tensor.t(), orientations: Nx.Tensor.t()}
Frame representation for orientation-aware FABRIK.
:positions- Joint positions as{n+1, 3}tensor:orientations- Joint orientations as{n+1, 4}tensor (WXYZ quaternions)
Functions
@spec backward_pass(Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t()) :: Nx.Tensor.t()
FABRIK backward reaching pass.
Pins the end effector to target, then walks toward the root placing each
joint at its segment length from the next. points is {n, 3}, lengths
{n - 1}. Returns the updated {n, 3} points.
@spec backward_pass_with_orientation( Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t(), number() ) :: {Nx.Tensor.t(), Nx.Tensor.t()}
Orientation-aware FABRIK backward reaching pass.
Pins the end-effector to target_position (and, when enforce_ori > 0.5, to
target_orientation), then walks toward the root placing each joint and
aligning its frame's Z-axis with the segment direction. positions is
{n, 3}, orientations {n, 4} (WXYZ), lengths {n - 1}. Returns
{positions, orientations}.
@spec fabrik( points :: Nx.Tensor.t(), lengths :: Nx.Tensor.t(), target :: Nx.Tensor.t(), max_iterations :: pos_integer(), tolerance :: float() ) :: {:ok, Nx.Tensor.t(), map()} | {:error, :unreachable | :max_iterations, map()}
Run the FABRIK algorithm on a chain of points.
Parameters
points- Nx tensor of shape{n+1, 3}representing joint/link positionslengths- Nx tensor of shape{n}representing segment lengthstarget- Nx tensor of shape{3}representing target positionmax_iterations- Maximum number of iterationstolerance- Convergence tolerance (distance to target)
Returns
{:ok, new_points, meta}- Converged successfully{:error, :unreachable, meta}- Target is beyond reach{:error, :max_iterations, meta}- Did not converge within max iterations
In all cases, meta contains:
:points- Final point positions:iterations- Number of iterations performed:residual- Final distance to target
@spec fabrik_with_orientation( frames :: frames(), lengths :: Nx.Tensor.t(), target_position :: Nx.Tensor.t(), target_orientation :: BB.Math.Quaternion.t() | nil, max_iterations :: pos_integer(), tolerance :: float(), opts :: keyword() ) :: {:ok, frames(), map()} | {:error, :unreachable | :max_iterations, map()}
Run FABRIK with full orientation support.
Parameters
frames- Map with:positions{n+1, 3}and:orientations{n+1, 4}tensorslengths- Segment lengths as{n}tensortarget_position- Target position as{3}tensortarget_orientation- Target quaternion asQuaternion.t()ornilfor position-onlymax_iterations- Maximum solver iterationstolerance- Position convergence tolerance (metres)opts- Options including:orientation_tolerance(radians, default 0.01)
Returns
{:ok, frames, meta}- Converged successfully{:error, :unreachable, meta}- Target beyond reach{:error, :max_iterations, meta}- Did not converge
Meta includes:
:frames- Final frame state:iterations- Iterations performed:residual- Position residual (metres):orientation_residual- Orientation residual (radians) ornil
@spec forward_pass(Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t()) :: Nx.Tensor.t()
FABRIK forward reaching pass.
Pins the root to root, then walks toward the end effector placing each joint
at its segment length from the previous. points is {n, 3}, lengths
{n - 1}. Returns the updated {n, 3} points.
@spec forward_pass_with_orientation( Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t() ) :: {Nx.Tensor.t(), Nx.Tensor.t()}
Orientation-aware FABRIK forward reaching pass.
Pins the root to root_position/root_orientation, then walks toward the
end-effector placing each joint and propagating orientation from its parent
along the segment direction. positions is {n, 3}, orientations {n, 4},
lengths {n - 1}. Returns {positions, orientations}.
@spec frames_to_points(frames()) :: Nx.Tensor.t()
Extract positions tensor from frames.
Place a point at desired_distance from anchor, along the direction from
anchor toward point_to_move. The per-joint reaching step shared by both
passes; a defn so it composes into them and is reusable on its own.
@spec points_to_frames(Nx.Tensor.t()) :: frames()
Convert points tensor to frames with identity orientations.
Useful for using the orientation-aware algorithm with position-only data.
@spec solve(Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t(), pos_integer(), float()) :: {Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t()}
Run the FABRIK iteration loop until convergence or max_iterations.
Composable defn entry point. Returns {points, iterations, residual} as
tensors — fabrik/5 wraps this with the reachability check and {:ok, ...}
classification. Each iteration runs backward_pass/3 then forward_pass/3.
Built from defn so it can be composed into larger numerical pipelines and
vectorised over a leading batch axis (e.g. solving many legs of a gait at
once).
@spec solve_with_orientation( Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t(), number(), pos_integer(), float(), float() ) :: {Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t()}
Run the orientation-aware FABRIK loop until convergence or max_iterations.
Composable defn companion to solve/5 that also tracks a {n, 4} quaternion
per joint. enforce_ori (1.0/0.0) selects whether the end-effector
orientation is pinned to target_orientation and whether orientation
convergence is required. Returns
{positions, orientations, iterations, residual, orientation_residual} as
tensors; fabrik_with_orientation/7 wraps this with reachability and {:ok, ...} classification.
@spec stretch_with_orientation( Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t() ) :: {Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t()}
Stretch the chain straight toward an unreachable target.
Lays joints along the root→target direction at their segment lengths,
propagating orientation along that direction. Returns
{positions, orientations, orientation_residual} (the residual is the angular
distance of the end-effector orientation from target_orientation).