Physically-Grounded Latent Action Learning for Contact-Rich Robot Manipulation

ECE 228, Spring 2026 (UC San Diego). Final project, Track 1 (Reproduce-and-Improve).

Author: Haoyang Li (hal212@ucsd.edu) Department: Computer Science and Engineering, UC San Diego

Overview

This repository contains the code accompanying the ECE 228 final project. The project reimplements the LAPO-style latent action model and proposes a simple change to its training objective: replace future-frame reconstruction with a symmetric InfoNCE alignment between the visual latent and an embedding of physical signals (tactile readings and proprioceptive state) recorded alongside the video. Physical signals are used only during training, so the inverse dynamics model remains visual-only at deployment time.

On the Touch in the Wild dataset (11 manipulation tasks, ~2.2M synchronized frames), the proposed objective lifts the action-probe coefficient of determination from the 0.005–0.026 basin of the reconstruction family to 0.300–0.464 across four contrastive variants. The decisive comparison is C1 (tactile-only anchor) at R² = 0.300 versus B1 (vision+tactile reconstruction) at R² = 0.026: the same physical modality is given in both cases, but only the contrastive objective recovers action-relevant information from it.

Headline result

A shuffled-tactile control (B2) in the reconstruction baseline produces essentially identical downstream R² to the unshuffled counterpart, showing that the reconstruction objective does not in fact use the tactile information it is given. The training objective, not the modality, is the limiting factor.

Repository layout

configs/
  train/              # YAML configs for B0–B7 and C1–C4
  train_3d/           # Configs for follow-up backbone experiments
  model/, dataset/    # Component-level configs
src/
  datasets/           # TouchWildDataset and friends
  models/             # VisualEncoder, TactileEncoder, LatentActionModel
  losses/             # Symmetric InfoNCE
  trainers/           # Training loop with KL + contact head
  utils/              # IO, metrics, visualisation helpers
scripts/
  train/              # Training entry points
  eval/               # Probe evaluation and policy fitting
  extract/            # Frozen backbone feature extraction
  preprocess/         # Dataset preprocessing utilities
  visualize/          # Latent-space and dataset visualisation
  round13/, round14/, round15/  # Follow-up experiments beyond the report scope
requirements.txt
setup_env.sh

Setup

# Create the conda environment
bash setup_env.sh

# Or, with conda directly
conda create -n contact-lat python=3.10 -y
conda activate contact-lat
pip install -r requirements.txt

Quick start

The full Touch in the Wild dataset is large (~1.4 TB) and not bundled with this repo. Download the dataset from the official source (https://touch-in-the-wild.github.io/) and place the .zarr.zip archives under data/raw/touch_wild/indoor_data/<task>/.

# Train the proposed C1 (tactile-only contrastive) model
python scripts/train/train_baseline.py --config configs/train/c1_tactile_contrastive.yaml

# Train the LAPO reconstruction baseline (B0)
python scripts/train/train_baseline.py --config configs/train/b0_visual_only.yaml

# Run the action-probe evaluation on a trained checkpoint
python scripts/eval/eval_policy_multitask.py \
    --checkpoint outputs/checkpoints/c1_tactile_contrastive/best.pt

Datasets

The primary dataset is Touch in the Wild (Zhu, Huang, Li, 2025), restricted to the eleven-task subset with full contact annotations. The follow-up scripts under scripts/extract/, scripts/round13/, etc. additionally consume DROID and HOT3D for the backbone-comparison work mentioned briefly in the report; those evaluations are outside the scope of the core ECE 228 deliverable.

Method (summary)

The model has the standard inverse/forward dynamics structure used by LAPO and Genie. A ResNet-18 encoder produces per-window features, an MLP inverse dynamics model maps the concatenated features to a 16-dim latent z, and a forward dynamics MLP predicts the next-window features from z and the history. Training optimises a symmetric InfoNCE loss between z and an embedding of physical signals from a separate encoder, with a small KL regulariser on the latent and an auxiliary contact-prediction head to prevent posterior collapse. At inference, the inverse dynamics model consumes only RGB.

See the project report for full method details, related work, experimental conditions, and limitations.

References

The implementation builds on LAPO (Schmidt and Jiang, 2024), Genie (Bruce et al., 2024), CLIP (Radford et al., 2021), and the Touch in the Wild dataset (Zhu, Huang, Li, 2025). Full citations are listed in the project report.

License

Code in this repository is released for educational use in connection with the ECE 228 course project. Third-party dependencies are governed by their own licenses.