penji2803/lidar-vjepa-turtlebot3

LiDAR V-JEPA (TurtleBot3, 360° 2-D lidar)

A small V-JEPA-style self-supervised encoder for sequences of 2-D lidar scans, pre-trained on 1.5M simulated TurtleBot3 frames with moving obstacles. Given the last 4 s of a 360-ray lidar (20 scans at 5 Hz), it produces token embeddings (10 time steps × 36 angular sectors × 128) or a pooled 1024-D last8 embedding (8 direction bins × 128) that encodes obstacle layout, nearby-obstacle motion and short-term collision risk.

This repo contains the full pre-trained JEPA (tokenizer, context encoder, EMA target encoder, action-conditioned predictor) and a standalone frozen encoder that needs only torch, numpy and safetensors.

File	What it is
encoder_last8.safetensors + encoder_config.json + jepa_encoder.py	Frozen target encoder → last8 embedding (1024-D, standardised). No transformers needed.
model.safetensors + config.json + lidar_jepa/	Full JEPA (all four parts, 3.2 M parameters) for fine-tuning, other poolings or continued pre-training. Needs transformers>=5.
example.py	Runs both.
figures/	Training curves and probe results.

Quick start

pip install torch numpy safetensors huggingface_hub   # + "transformers>=5,<6" for the full model

from huggingface_hub import snapshot_download
import sys, numpy as np
path = snapshot_download('penji2803/lidar-vjepa-turtlebot3')
sys.path.insert(0, path)

from jepa_encoder import JepaLast8Encoder
enc = JepaLast8Encoder.load(path)                 # CPU is fine: ~5 ms per window
scans = np.full((20, 360), 3.5, np.float32)       # last 20 scans, oldest first, metres
valid = np.ones(20, bool)
z = enc.embed(scans, valid)                       # (1024,) = 8 directions x 128

from lidar_jepa.hub import load_jepa              # full model (needs transformers)
model = load_jepa(path)

Input format

• scans: (20, 360) float32, metres, the last 20 scans at 5 Hz, oldest first. Ranges are clipped to [0.12, 3.5] m, and 3.5 means "no return" (map inf to 3.5).
• Ray i is at angle i · 2π / 360: ray 0 points straight ahead, and angles increase counter-clockwise (ROS LaserScan convention, as on the TurtleBot3 LDS).
• valid: (20,) bool. At the start of an episode, left-pad with copies of the first scan and mark those frames False. Padded frames are masked out of attention.

Output (last8)

The 1024 numbers are 8 blocks of 128, one per direction bin: bin 0 is straight ahead (−20°…+20°), and bins go counter-clockwise (bin 2 = left, 4 = behind, 6 = right). Each block is the mean of the most recent time step's sector tokens in that direction, standardised with training-set statistics. FrozenLidarEncoder(..., pool=...) in lidar_jepa/inference.py offers other poolings: meanmax, sector8, last_t and flat.

Model

scans (20, 360) ─► 2 channels: inverse range (0.12/r, 0 = no return) + validity mask
                ─► tubelets 2 frames × 10 rays (Conv2d 2→128)  ─► 10 × 36 = 360 tokens × 128
                ─► LayerNorm + circular sin-cos position encoding (time × angle; sector 35 neighbours sector 0)
                ─► ViT encoder: 6 pre-LN layers, 4 heads, width 128, MLP 512, GELU

Part	Layers	Heads	Width	Params	Role
Context encoder	6	4	128	1.19 M	Encodes the visible tokens
Target encoder (EMA)	6	4	128	1.19 M	Encodes all tokens; this is the inference encoder
Predictor	4	4	128	0.84 M	Predicts hidden tokens' target embeddings from the context, one action token per tubelet step (velocity commands) and mask tokens

The encoders are Hugging Face ViTModel stacks run on the lidar tokens (the patch embedding is not used). A dependency-free re-implementation of the target branch is in jepa_encoder.py; it matches the full model to within 1e-5.

Self-supervised objective (V-JEPA style).

• Masks: sector tubes, i.e. the same angular sectors hidden through all 10 time steps. Two mask families are used in every batch:

  • short-range: 8 blocks of 3 sectors (30°);
  • long-range: 2 blocks of 12 sectors (120°).

  Each hides 24 of the 36 sectors (240 target tokens, 120 context tokens).

• Loss: L1 between the predicted embeddings and the layer-normed target-encoder embeddings.

• Target encoder: an EMA of the context encoder.

Training

Data	penji2803/turtlebot3-moving-obstacles-lidar-1.5m: 510k windows of 20 frames (from 1.5M rows / 14k episodes). Validation: the dataset's independent 150k-row set (51k windows, no shared episodes).
Simulator	Gazebo, ROBOTIS TurtleBot3 Burger, Stage 3 arena with 4 moving obstacles; clean sensors
Optimiser	AdamW, lr 5e-4 (warm-up 5 % + cosine to 1e-6), weight decay 0.04 → 0.4, grad clip 1.0, batch 128, bf16
EMA	momentum 0.996 → 1.0
Schedule	40 epochs on one H100; this checkpoint is epoch 16 (best validation loss)

Results:

• Validation loss: 0.168, against 0.74 for the trivial "predict zeros" baseline (−77 %). Training loss at the same epoch was 0.171, so there is no overfitting.

• Both mask families were learned:

  Mask family	Validation loss, epoch 1 → 16
  Short-range	0.327 → 0.164
  Long-range	0.358 → 0.171

• No representation collapse: the target-embedding variance fraction stayed between 0.58 and 0.86.

What the embedding contains. These are frozen last8 features with a small MLP probe on top, trained on 60k windows and tested on the 51k validation windows:

Probe target	last8 (1024)	Random-init encoder (1024)	Raw scans, 20 × 360
Nearest-obstacle distance (R²)	0.924	0.695	0.694
Min range per 10° sector, now (R²)	0.805	0.704	0.867
Min range 1 s ahead (R²)	0.466	0.368	0.317
Collision within 1 s (AUROC)	0.993	0.969	0.935
Last velocity command (accuracy)	0.699	0.567	0.581

A linear probe reaches 0.899 / 0.786 / 0.453 / 0.968 / 0.574 on the same five targets.

Intended use and limitations

• Intended use: a compact lidar representation for mobile-robot research: RL observations, collision-risk or clearance prediction heads, world-model experiments, and fine-tuning on related 2-D lidar data.
• Simulation only: trained purely in Gazebo (one arena, cylindrical moving obstacles, clean sensors). Expect a domain gap on real robots and other environments; fine-tuning is advised.
• Sensor format: it expects a 360-ray, 360°, 5 Hz lidar with 3.5 m range. Other sensors need resampling, or retraining the tokenizer.
• Noise sensitivity: the embedding is fairly sensitive to sensor noise. Gaussian range noise with σ = 0.01 m plus 1–2 % ray dropout moves the last8 embedding about 0.42× as much as one real time step does; raw scans move about 0.03×. Downstream heads trained on noise-augmented inputs help.
• Feeding it straight into RL may not help: in our TurtleBot3 experiments, feeding the raw 1024-D embedding directly to PPO did not beat a 74-input lidar baseline. Reducing it first with small supervised heads (per-direction 1 s collision risk and clearance) gave a policy on par with or slightly better than the baseline.
• Not a safety-certified component.

Code

lidar_jepa/ contains the model code used for pre-training:

• preprocess, tokenizer, masking, encoders, predictor, model: the architecture;
• inference: poolings and FrozenLidarEncoder;
• train: training loop and schedules;
• hub: loading these weights.

train.py / data.py read a frame store built from the dataset with windows.py (set LIDAR_JEPA_DATA). They are included as a reference implementation and are not a turnkey training script.

License and citation

MIT. The design follows V-JEPA (Bardes et al., 2024, Revisiting Feature Prediction for Learning Visual Representations from Video), adapted from video patches to 2-D lidar tubelets.

@misc{lidar_vjepa_turtlebot3,
  title  = {LiDAR V-JEPA: self-supervised predictive representations for 2-D lidar (TurtleBot3)},
  author = {penji2803},
  year   = {2026},
  url    = {https://huggingface.co/penji2803/lidar-vjepa-turtlebot3}
}