DynaMo: In-Domain Dynamics Pretraining for Visuo-Motor Control

Getting started

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

• 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

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_*.

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.