| --- |
| 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). |
|
|
