README.md

---
license: mit
library_name: jax
tags:
- robotics
- motion-planning
- datafarm
- franka
- numpy
datasets:
- lerobot/droid_1.0.1
---

# 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](https://github.com/abhaybd/DataFarm-PyRoki).

Original method: [DATAFARM: Distribution-Aligned Task and Motion Planning for
Fine-Tuning Vision-Language-Action Models](https://arxiv.org/abs/2609.12316)
by Samrat Sahoo, Yixuan Huang and Tom Silver (2026).
See the [original DATAFARM repository](https://github.com/SamratSahoo/DATAFARM)
and [project website](https://prpl-group.com/datafarm/).

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:

```sh
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.

```python
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](https://github.com/SamratSahoo/DATAFARM/tree/fd403ea6eb16163261a13439fe9bbc10a980f50d).

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](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`:

```text
ed6296980f448704c92e88cc91e9767a7dd4b758e82b93b80adc80178ee825a9
```

## License and attribution

The original encoder and weights are MIT-licensed, Copyright (c) 2026 Samrat Sahoo.
See [DATAFARM-LICENSE.txt](DATAFARM-LICENSE.txt). The JAX port is maintained in
[DataFarm-PyRoki](https://github.com/abhaybd/DataFarm-PyRoki).