abhaybd/DataFarm-Jax-PosturePrior-Franka

DATAFARM JAX Posture Prior — Franka

Pretrained full-covariance Gaussian-mixture parameters for scoring Franka arm joint configurations in JAX. The implementation and PyRoki integration are in DataFarm-PyRoki.

Original method: DATAFARM: Distribution-Aligned Task and Motion Planning for Fine-Tuning Vision-Language-Action Models by Samrat Sahoo, Yixuan Huang and Tom Silver (2026). See the original DATAFARM repository and project website.

Use the negative log density (NLL) to rank already feasible IK solutions. Lower NLL favors configurations closer to the reference posture distribution. This is a prior over individual joint configurations. The separate motion encoder scores trajectories using a learned motion embedding.

License

The parameter artifact retains the upstream NVIDIA License, included verbatim in LICENSE and CUTAMP-LICENSE.txt. Section 3.3 limits use to non-commercial research or evaluation, with an explicit exception for NVIDIA Corporation and its affiliates. Redistribution must preserve the license and upstream notices. The complete license governs use of these parameters; they are not relicensed under the code repository's MIT license.

Model contract

• Reference embodiment: fr3_franka in the upstream DROID posture metadata.
• Mixture: 16 components with full covariance over seven arm coordinates.
• Joint units: radians; order is joint1 through joint7, proximal to distal.
• Gripper coordinates are excluded.
• Reference data, as recorded upstream: 4,000 episodes and 987,631 non-idle frames.
• Serialization: droid_joint_prior.npz, schema version 1, pickle-free NumPy arrays.

The archive contains mixture weights, means, precision factors, training location and scale, and explicit model/joint identity. For standardized coordinates z = (q - location) / scale, each factor satisfies precision = C @ C.T. The JAX score includes Gaussian normalization and the joint-unit change-of-variables term. These constant terms preserve the upstream candidate ordering while reporting a full NLL.

Load and score in JAX

uv venv
uv pip install "datafarm-pyroki @ git+https://github.com/abhaybd/DataFarm-PyRoki.git@d465221b9cf88efbb8ecd53fb10f04291dd97e7a" huggingface_hub

This model is public and downloads without authentication.

import jax
import jax.numpy as jnp
from huggingface_hub import hf_hub_download
from datafarm_pyroki import load_joint_prior

checkpoint = hf_hub_download(
    repo_id="abhaybd/DataFarm-Jax-PosturePrior-Franka",
    filename="droid_joint_prior.npz",
)
prior = load_joint_prior(checkpoint)

# Synthetic configuration in training joint order; replace with your candidate.
q = jnp.array([0.0, -0.5, 0.0, -2.0, 0.0, 1.5, 0.8])
nll = jax.jit(prior.nll)(q)
gradient = jax.grad(prior.nll)(q)
# prior.nll also supports [..., 7] batches.

For IK selection, construct an explicit joint mapping and provide your IK and collision feasibility results:

from datafarm_pyroki import JointMap, select_ik_solution

# `robot_joint_names` and `ik_candidates` come from your robot/IK solver.
# Supply aliases={model_joint_name: robot_joint_name} when URDF names differ.
joint_map = JointMap.from_names(robot_joint_names, prior.joint_names)
selected = select_ik_solution(
    prior,
    ik_candidates,  # [..., candidates, robot_joints]
    joint_map=joint_map,
    feasible=ik_success_and_collision_mask,  # [..., candidates]
    lower_limits=robot_joint_lower_limits,
    upper_limits=robot_joint_upper_limits,
)
q_goal = selected.positions  # complete configuration in robot joint order

rank_ik_solutions is the JIT-compatible ranking API; handle its all-invalid mask explicitly. select_ik_solution is a host-side convenience wrapper and raises NoFeasibleIKSolutions when a query has no valid candidate. Ranking is stable for tied scores. Coordinates outside the model mapping are retained and checked for finite values and supplied joint limits.

Inference and ranking need JAX and NumPy. PyTorch, sklearn, cuTAMP and PyRoki are not required to load or evaluate this prior. An NLL score does not check pose, collision or task feasibility and does not establish compatibility with another embodiment.

Provenance and validation

Parameters are extracted losslessly, without refitting, from the posture_ref.npz used by DATAFARM's pinned cuTAMP revision. The original source URL, revision, data provenance and source/converted checksums are retained in joint_prior_provenance.json.

The implementation tests NLL against an independent Gaussian-mixture calculation, checks gradients and serialization, and verifies batched candidate ranking with joint reordering, feasibility masks, nonfinite values and joint limits. These checks establish numerical behavior, not physical-robot performance.

Files and integrity

• droid_joint_prior.npz: mixture parameters and training coordinate metadata.
• joint_prior_provenance.json: upstream source and conversion checksums.
• LICENSE and CUTAMP-LICENSE.txt: verbatim upstream NVIDIA License.

SHA256 of droid_joint_prior.npz:

3252f89dfb3a675d726ff6ef5fee3371244da6969d11a5bb7611425bebcf54e0