Upload pretrained DATAFARM Franka motion encoder for JAX

DATAFARM-LICENSE.txt

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Samrat Sahoo
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,3 +1,124 @@
 ---
 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).
+
+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 stroke 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
+```
+
+For a private repository, authenticate with `hf auth login` before downloading.
+
+```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 stroke 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).

droid_encoder.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ed6296980f448704c92e88cc91e9767a7dd4b758e82b93b80adc80178ee825a9
+size 1331772

provenance.json

@@ -0,0 +1,46 @@
+{
+  "model_name": "droid-franka",
+  "joint_names": [
+    "joint1",
+    "joint2",
+    "joint3",
+    "joint4",
+    "joint5",
+    "joint6",
+    "joint7"
+  ],
+  "sample_rate_hz": 15.0,
+  "source_repository": "https://github.com/SamratSahoo/DATAFARM",
+  "source_commit": "fd403ea6eb16163261a13439fe9bbc10a980f50d",
+  "source_checkpoint": "https://raw.githubusercontent.com/SamratSahoo/DATAFARM/fd403ea6eb16163261a13439fe9bbc10a980f50d/encoder/checkpoints/encoder.pt",
+  "source_checkpoint_sha256": "e31ef96b6ffc0af8ce890e8774b9b57d3fa6e91c48ba9125d1e3574219c30c5b",
+  "converted_checkpoint_sha256": "ed6296980f448704c92e88cc91e9767a7dd4b758e82b93b80adc80178ee825a9",
+  "source_model": "https://github.com/SamratSahoo/DATAFARM/blob/fd403ea6eb16163261a13439fe9bbc10a980f50d/encoder/model.py",
+  "license": "MIT; Copyright (c) 2026 Samrat Sahoo; see DATAFARM-LICENSE.txt",
+  "conversion": "Lossless extraction of state_dict tensors and reference statistics to NPZ",
+  "training_metadata": {
+    "dataset": "lerobot/droid_1.0.1",
+    "dataset_codebase_version": "v3.0",
+    "train_protocol": "full DROID lerobot/droid_1.0.1 (v3.0); fingerprint regression + Fisher discriminant (within/between latent scatter) on balanced DROID/cuTAMP batches; encoder on a DROID subsample + cuTAMP, val-selected by held-out pair-silhouette; cluster stats over ALL episodes",
+    "n_droid_train": 30000,
+    "n_droid_stats": 94774,
+    "stats_max_segments": 120000,
+    "hparams": {
+      "d": 16,
+      "beta_max": 0.005,
+      "dropout": 0.3,
+      "disc_w": 1.0,
+      "batch": 128,
+      "n_crops_d": 6,
+      "n_crops_c": 16,
+      "balanced": true
+    },
+    "metrics": {
+      "val_pair_silhouette": 0.8552306294441223,
+      "val_probe": 1.0,
+      "selected_epoch": 50
+    },
+    "kl_droid_mean": 17.193614959716797,
+    "maha2_droid_mean": 7.039307359411732
+  }
+}