/key_dynam

Code for paper "Keypoints into the Future: Self-Supervised Correspondence in Model-Based RL"

Primary LanguagePython

Code Overview

Our approach requires

  1. Generating a dataset
  2. Training the visual models (dense correspondence, transporter, etc.)
  3. Training the dynamics models
  4. Evaluating closed-loop MPC performance

Each experiment/task has it's own folder, e.g. key_dynam/experiments/drake_pusher_slider. This folder contains scripts that perform each of the four steps listed above.

Environment setup

Our code is setup to run inside a docker environment. To build the docker image use

cd key_dynam/docker && ./docker_build.py

and to run the docker container use

cd key_dynam/docker && ./docker_run.py

If setting up the env using conda, first go to dense object net repo to install requirements. Then update requirements here. Note: Python=3.6

Install bazel: following https://docs.bazel.build/versions/main/install-ubuntu.html Install drake: following https://drake.mit.edu/python_bindings.html#installation and download v0.25.0

Generating a dataset

To generate a dataset use for example key_dynam/experiments/drake_pusher_slider/collect_episodes.py

Training visual model

Use the train_dense_descriptor_vision function of the key_dynam/experiments/drake_pusher_slider/train_and_evaluate.py file.

Training and Evaluating the dynamics model

Use the methods DD_3D, GT_3D, transporter_3D, train_conv_autoencoder_dynamics in the key_dynam/experiments/drake_pusher_slider/train_and_evaluate.py file.