first commit

README.md

@@ -0,0 +1,90 @@
+# ACT: Action Chunking with Transformers
+
+### Project Website: https://tonyzhaozh.github.io/aloha/
+
+This repo contains the implementation of ACT, together with 2 simulated environments:
+Transfer Cube and Bimanual Insertion. You can train and evaluate ACT in sim (tested) or real (ongoing).
+
+
+### Repo Structure
+- ``imitate_episodes.py`` Train and Evaluate ACT
+- ``policy.py`` An adaptor for ACT policy
+- ``detr`` Model definitions of ACT, modified from DETR
+- ``sim_env.py`` Mujoco + DM_Control environments with joint space control
+- ``ee_sim_env.py`` Mujoco + DM_Control environments with EE space control
+- ``scripted_policy.py`` Scripted policies for sim environments
+- ``constants.py`` Constants shared across files
+- ``utils.py`` Utils such as data loading and helper functions
+- ``visualize_episodes.py`` Save videos from a .hdf5 dataset
+
+
+### Installation
+
+    conda create -n aloha python=3.8
+    conda activate aloha
+    pip install torchvision
+    pip install torch
+    pip install pyquaternion
+    pip install pyyaml
+    pip install rospkg
+    pip install pexpect
+    pip install mujoco
+    pip install dm_control
+    pip install opencv-python
+    pip install matplotlib
+    pip install einops
+    pip install packaging
+    pip install h5py
+    pip install h5py_cache
+    cd act/detr && pip install -e .
+
+### Example Usages
+
+To set up a new terminal, run:
+
+    conda activate aloha
+    cd <path to act repo>
+
+### Simulated experiments
+
+We use ``transfer_cube`` task in the examples below. Another option is ``insertion``.
+To generated 50 episodes of scripted data, run:
+
+    python3 record_sim_episodes.py \
+    --task_name transfer_cube \
+    --dataset_dir <data save dir> \
+    --num_episodes 50
+
+To can add the flag ``--onscreen_render`` to see real-time rendering.
+To visualize the episode after it is collected, run
+
+    python3 visualize_episodes.py --dataset_dir <data save dir> --episode_idx 0
+
+To train ACT:
+    
+    # Transfer Cube task
+    python3 imitate_episodes.py \
+    --dataset_dir <data save dir> \
+    --ckpt_dir <ckpt dir> \
+    --policy_class ACT --kl_weight 10 --chunk_size 100 --hidden_dim 256 --batch_size 8 --dim_feedforward 2048 \
+    --task_name transfer_cube --seed 0 \
+    --temporal_agg \
+    --num_epochs 1000  --lr 1e-4
+    
+    # Bimanual Insertion task
+    python3 imitate_episodes.py \
+    --dataset_dir <data save dir> \
+    --ckpt_dir <ckpt dir> \
+    --policy_class ACT --kl_weight 10 --chunk_size 100 --hidden_dim 256 --batch_size 8 --dim_feedforward 2048 \
+    --task_name insertion --seed 0 \
+    --temporal_agg \
+    --num_epochs 2000  --lr 1e-5
+
+To evaluate the policy, run the same command but add ``--eval``. The success rate
+should be around 85% for transfer cube, and around 50% for insertion.
+Videos will be saved to ``<ckpt_dir>`` for each rollout.
+You can also add ``--onscreen_render`` to see real-time rendering during evaluation.
+
+
+
+