README.md

# **DynaMo**: In-Domain Dynamics Pretraining for Visuo-Motor Control
[[Paper]](https://arxiv.org/abs/2409.12192) [[Project Website]](https://dynamo-ssl.github.io/) [[Data]](https://osf.io/kxehw/)

[Zichen Jeff Cui](https://jeffcui.com/), [Hengkai Pan](https://www.ri.cmu.edu/ri-people/hengkai-pan/), [Aadhithya Iyer](https://aadhithya14.github.io/), [Siddhant Haldar](https://siddhanthaldar.github.io/) and [Lerrel Pinto](https://www.lerrelpinto.com/), New York University

This repo contains code for DynaMo visual pretraining, and for reproducing sim environment experiments. Datasets will be uploaded soon.


## Getting started
The following assumes our current working directory is the root directory of this project repo; tested on Ubuntu 22.04 LTS (amd64).
### Setting up the project environments
- Install the project environment:
  ```
  conda env create --file=conda_env.yml
  ```
- Activate the environment:
  ```
  conda activate dynamo-repro
  ```
- To enable logging, log in with a `wandb` account:
  ```
  wandb login
  ```
  Alternatively, to disable logging altogether, set the environment variable `WANDB_MODE`:
  ```
  export WANDB_MODE=disabled
  ```

### Getting the training datasets
[Get the dataset here](https://osf.io/kxehw/).

(Updated Sep 29: sim kitchen dataset now supports lazy loading: set `prefetch=False` in the sim kitchen configs. If you encounter errors, try downloading the latest dataset zips from the link above.)
- Download all files in the `datasets` directory, combine all partitions, and unzip:
  ```
  zip -s- dynamo_repro_datasets.zip -O combined.zip
  unzip combined.zip
  ```
- In `./configs/env_vars/env_vars.yaml`, set `dataset_root` to the unzipped parent directory containing all datasets.
- In `./eval_configs/env_vars/env_vars.yaml`, set `dataset_root` to the unzipped parent directory containing all datasets.
- In `./eval_configs/env_vars/env_vars.yaml`, set `save_path` to where you want to save the rollout results (e.g. root directory of this repo).
- Environments:
  - `sim_kitchen`: Franka kitchen environment
  - `block_push_multiview`: Block push environment
  - `libero_goal`: LIBERO Goal environment
  - `pusht`: Push-T environment

## Reproducing experiments
The following assumes our current working directory is the root directory of this project repo.

To reproduce the experiment results, the overall steps are:
1. Activate the conda environment with
   ```
   conda activate dynamo-repro
   ```

2. Train the visual encoder with `python3 train.py --config-name=train_*`. A model snapshot will be saved to `./exp_local/...`;
3. In `eval_configs/encoder`, in the corresponding environment config, set the encoder file path `f` to the saved snapshot;
4. Eval with `python3 online_eval.py --config-name=train_*`.

See below for detailed steps for each environment.


### Franka Kitchen
- Train the encoder:
  ```
  python3 train.py --config-name=train_sim_kitchen
  ```
  Snapshots will be saved to a new timestamped directory `./exp_local/{date}/{time}_train_sim_kitchen_dynamo`.

  The encoder snapshot will be at `./exp_local/{date}/{time}_train_sim_kitchen_dynamo/encoder.pt`.
- In `eval_configs/encoder/kitchen_dynamo.yaml`, set `SNAPSHOT_PATH` to the absolute path of the encoder snapshot above.
- Evaluation:
  ```
  MUJOCO_GL=egl python3 online_eval.py --config-name=train_sim_kitchen
  ```

### Block Pushing
- Train the encoder:
  ```
  python3 train.py --config-name=train_blockpush
  ```
  Snapshots will be saved to a new timestamped directory `./exp_local/{date}/{time}_train_blockpush_dynamo`.

  The encoder snapshot will be at `./exp_local/{date}/{time}_train_blockpush_dynamo/encoder.pt`.
- In `eval_configs/encoder/blockpush_dynamo.yaml`, set `SNAPSHOT_PATH` to the absolute path of the encoder snapshot above.
- Evaluation:
  ```
  ASSET_PATH=$(pwd) python3 online_eval.py --config-name=train_blockpush
  ```
  (Evaluation requires including this repository in `ASSET_PATH`.)

### Push-T
- Train:
  ```
  python3 train.py --config-name=train_pusht
  ```
  Snapshots will be saved to a new timestamped directory `./exp_local/{date}/{time}_train_pusht_dynamo`.
  
  The encoder snapshot will be at `./exp_local/{date}/{time}_train_pusht_dynamo/encoder.pt`
- In `eval_configs/encoder/pusht_dynamo.yaml`, set `SNAPSHOT_PATH` to the absolute path of the encoder snapshot above.
- Evaluation:
  ```
  python3 online_eval.py --config-name=train_pusht
  ```

### LIBERO Goal
- Train:
  ```
  python3 train.py --config-name=train_libero_goal
  ```
  Snapshots will be saved to a new timestamped directory `./exp_local/{date}/{time}_train_libero_goal_dynamo`.

  The encoder snapshot will be at `./exp_local/{date}/{time}_train_libero_goal_dynamo/encoder.pt`
- In `eval_configs/encoder/libero_dynamo.yaml`, set `SNAPSHOT_PATH` to the absolute path of the encoder snapshot above.
- Evaluation:
  ```
  MUJOCO_GL=egl python3 online_eval.py --config-name=train_libero_goal
  ```

## Train on your own dataset
- Plug in your dataset in these files:
  - `datasets/your_dataset.py`
  - `configs/env/your_dataset.yaml`
  - `configs/env_vars/env_vars.yaml`

- Check the inverse/forward model configs:
  - `configs/train_your_dataset.yaml`
    - This is the main config.
  - `configs/ssl/dynamo_your_dataset.yaml`
    - If the model converges slowly, try setting `ema_beta` to `null` to use SimSiam instead of EMA encoder during training.
  - `configs/projector/inverse_dynamics_your_dataset.yaml`
    - We find that setting the inverse dynamics `output_dim` to approximately the underlying state dimension usually works well.
      - For sim environments, this is the state-based observation dimension.
      - For real environments, e.g. a 7DoF robot arm + gripper (1D) manipulating a rigid object (6D), this would be ~16 dimensions.

- Add linear probes for training diagnostics:
  - `workspaces/your_workspace.py`
    - This template computes linear probe and nearest neighbor MSE from the image embeddings to states/actions, for monitoring training convergence.
    - It assumes that your dataset class has `states` (`batch` x `time` x `state_dim`) and `actions` (`batch` x `time` x `action_dim`) attributes.
      - For a real-world dataset, you can use proprioception as the state.