abhaybd/DataFarm-Jax-MotionEncoder-Franka

DATAFARM JAX Motion Encoder — Franka

Pretrained DATAFARM motion encoder weights converted losslessly to a pickle-free NumPy archive for inference 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.

The frozen encoder maps a joint trajectory to a 16-dimensional latent mean. A fixed DROID reference Gaussian supplies the squared Mahalanobis motion-style cost used during trajectory optimization and subsequent trajectory segment retiming.

Model contract

• Embodiment: DROID Franka arm, seven joint-position coordinates in radians.
• Training order: joint1, joint2, joint3, joint4, joint5, joint6, joint7, proximal to distal. Gripper coordinates are excluded.
• Training sample rate: 15 Hz. The loader accepts a source timestep and resamples.
• Input: [time, 7] or [batch, time, 7]; internally standardized position, velocity, acceleration and jerk features.
• Output: 16-dimensional latent mean, with a saved reference mean and precision.
• Serialization: droid_encoder.npz, schema version 1; load with allow_pickle=False.

Coordinate names record training order. Supply an explicit JointMap for URDF aliases in PyRoki; a matching joint count alone does not establish compatibility with another robot.

Load in JAX

Install the tested implementation in a uv environment:

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_encoder

checkpoint = hf_hub_download(
    repo_id="abhaybd/DataFarm-Jax-MotionEncoder-Franka",
    filename="droid_encoder.npz",
)
encoder = load_encoder(checkpoint)

# Synthetic input, in training joint order; replace with your trajectory.
q = jnp.zeros((32, 7))
latent = encoder.latent(q, source_dt=0.15)
cost, gradient = jax.jit(
    jax.value_and_grad(lambda positions: encoder.cost(positions, source_dt=0.15))
)(q)

Inference uses JAX and does not require PyTorch, CUDA, cuRobo or TiPToP. The implementation supports CPU and appropriate JAX accelerator backends. See the code repository for whole-leg PyRoki costs and standalone trajectory segment retiming.

Architecture

Six dilated Conv1D filterbank branches contribute masked mean/std features. A parallel strided temporal CNN contributes masked mean/std/max features. The concatenated 672 features pass through an MLP with LayerNorm and exact GELU, then a linear head produces latent mean and log variance. The archive preserves all 77 state-dict tensors, normalization buffers, reference statistics and the auxiliary training head. Planning uses the deterministic latent mean.

Training provenance

This repository distributes a conversion of the released weights, without retraining. The original checkpoint is from SamratSahoo/DATAFARM.

The artifact's own metadata describes 30,000 DROID training episodes plus cuTAMP motion, fingerprint regression and a Fisher-discriminant objective on balanced DROID/cuTAMP batches, selected using held-out pair silhouette. Its reference latent statistics are DROID-derived. This metadata differs from the current upstream README/training script; provenance.json retains the checkpoint metadata and exact source/conversion checksums.

Validation and use

The JAX implementation is checked against the original PyTorch network for preprocessing, latent outputs, masked inference, motion costs and gradients. The code repository also tests float32/float64 batch consistency, retiming profile/duration gradients, joint mapping and PyRoki integration. Numerical parity does not establish performance on another embodiment or physical robot.

This is a learned motion-style score, not a collision or kinematic feasibility check. Keep your task, joint-limit and collision constraints when using it in a planner. The code repository documents the relevant validation hooks.

Files and integrity

• droid_encoder.npz: converted encoder weights and normalization/reference statistics.
• provenance.json: source revision, training metadata and SHA256 hashes.
• DATAFARM-LICENSE.txt: original MIT license and attribution.

SHA256 of droid_encoder.npz:

ed6296980f448704c92e88cc91e9767a7dd4b758e82b93b80adc80178ee825a9

License and attribution

The original encoder and weights are MIT-licensed, Copyright (c) 2026 Samrat Sahoo. See DATAFARM-LICENSE.txt. The JAX port is maintained in DataFarm-PyRoki.