search

PKU-MARL / DexterousHands

This is a library that provides dual dexterous hand manipulation tasks through Isaac Gym
https://pku-marl.github.io/DexterousHands/
Apache License 2.0
653 stars 78 forks source link
deep-reinforcement-learning dexterous-robotic-hand reinforcement-learning

readme

Bi-DexHands: Bimanual Dexterous Manipulation via Reinforcement Learning

PyPI Organization Unittest Docs GitHub license

Update

[2023/02/09] We re-package the Bi-DexHands. Now you can call the Bi-DexHands' environments not only on the command line, but also in your Python script. check our README Use Bi-DexHands in Python scripts below.

[2022/11/24] Now we support visual observation for all the tasks, check this document for visual input.

[2022/10/02] Now we support for the default IsaacGymEnvs RL library rl-games, check our README below.

Bi-DexHands (click bi-dexhands.ai) provides a collection of bimanual dexterous manipulations tasks and reinforcement learning algorithms. Reaching human-level sophistication of hand dexterity and bimanual coordination remains an open challenge for modern robotics researchers. To better help the community study this problem, Bi-DexHands are developed with the following key features:

Contents of this repo are as follows:

For more information about this work, please check our paper.

Installation

Details regarding installation of IsaacGym can be found here. We currently support the Preview Release 3/4 version of IsaacGym.

Pre-requisites

The code has been tested on Ubuntu 18.04/20.04 with Python 3.7/3.8. The minimum recommended NVIDIA driver version for Linux is 470.74 (dictated by support of IsaacGym).

It uses Anaconda to create virtual environments. To install Anaconda, follow instructions here.

Ensure that Isaac Gym works on your system by running one of the examples from the python/examples directory, like joint_monkey.py. Please follow troubleshooting steps described in the Isaac Gym Preview Release 3/4 install instructions if you have any trouble running the samples.

Once Isaac Gym is installed and samples work within your current python environment, install this repo:

Install from source code

You can also install this repo from the source code:

pip install -e .

Introduction

This repository contains complex dexterous hands control tasks. Bi-DexHands is built in the NVIDIA Isaac Gym with high performance guarantee for training RL algorithms. Our environments focus on applying model-free RL/MARL algorithms for bimanual dexterous manipulation, which are considered as a challenging task for traditional control methods.

Getting Started

Tasks

Source code for tasks can be found in envs/tasks. The detailed settings of state/action/reward are in here.

So far, we release the following tasks (with many more to come):

Environments Description Demo
ShadowHand Over These environments involve two fixed-position hands. The hand which starts with the object must find a way to hand it over to the second hand.
ShadowHandCatch Underarm These environments again have two hands, however now they have some additional degrees of freedom that allows them to translate/rotate their centre of masses within some constrained region.
ShadowHandCatch Over2Underarm This environment is made up of half ShadowHandCatchUnderarm and half ShadowHandCatchOverarm, the object needs to be thrown from the vertical hand to the palm-up hand
ShadowHandCatch Abreast This environment is similar to ShadowHandCatchUnderarm, the difference is that the two hands are changed from relative to side-by-side posture.
ShadowHandCatch TwoCatchUnderarm These environments involve coordination between the two hands so as to throw the two objects between hands (i.e. swapping them).
ShadowHandLift Underarm This environment requires grasping the pot handle with two hands and lifting the pot to the designated position
ShadowHandDoor OpenInward This environment requires the closed door to be opened, and the door can only be pulled inwards
ShadowHandDoor OpenOutward This environment requires a closed door to be opened and the door can only be pushed outwards
ShadowHandDoor CloseInward This environment requires the open door to be closed, and the door is initially open inwards
ShadowHand BottleCap This environment involves two hands and a bottle, we need to hold the bottle with one hand and open the bottle cap with the other hand
ShadowHandPush Block This environment requires both hands to touch the block and push it forward
ShadowHandOpen Scissors This environment requires both hands to cooperate to open the scissors
ShadowHandOpen PenCap This environment requires both hands to cooperate to open the pen cap
ShadowHandSwing Cup This environment requires two hands to hold the cup handle and rotate it 90 degrees
ShadowHandTurn Botton This environment requires both hands to press the button
ShadowHandGrasp AndPlace This environment has a bucket and an object, we need to put the object into the bucket

Training

Training Examples

RL/MARL Examples

For example, if you want to train a policy for the ShadowHandOver task by the PPO algorithm, run this line in bidexhands folder:

python train.py --task=ShadowHandOver --algo=ppo

To select an algorithm, pass --algo=ppo/mappo/happo/hatrpo/... as an argument. For example, if you want to use happo algorithm, run this line in bidexhands folder:

python train.py --task=ShadowHandOver --algo=happo

Supported Single-Agent RL algorithms are listed below:

Supported Multi-Agent RL algorithms are listed below:

Multi-task/Meta RL Examples

The training method of multi-task/meta RL is similar to the RL/MARL, it is only need to select the multi-task/meta categories and the corresponding algorithm. For example, if you want to train a policy for the ShadowHandMT4 categories by the MTPPO algorithm, run this line in bidexhands folder:

python train.py --task=ShadowHandMetaMT4 --algo=mtppo

Supported Multi-task RL algorithms are listed below:

Supported Meta RL algorithms are listed below:

Gym-Like API

We provide a Gym-Like API that allows us to get information from the Isaac Gym environment. Our single-agent Gym-Like wrapper is the code of the Isaac Gym team used, and we have developed a multi-agent Gym-Like wrapper based on it:

class MultiVecTaskPython(MultiVecTask):
    # Get environment state information
    def get_state(self):
        return torch.clamp(self.task.states_buf, -self.clip_obs, self.clip_obs).to(self.rl_device)

    def step(self, actions):
        # Stack all agent actions in order and enter them into the environment
        a_hand_actions = actions[0]
        for i in range(1, len(actions)):
            a_hand_actions = torch.hstack((a_hand_actions, actions[i]))
        actions = a_hand_actions
        # Clip the actions
        actions_tensor = torch.clamp(actions, -self.clip_actions, self.clip_actions)
        self.task.step(actions_tensor)
        # Obtain information in the environment and distinguish the observation of different agents by hand
        obs_buf = torch.clamp(self.task.obs_buf, -self.clip_obs, self.clip_obs).to(self.rl_device)
        hand_obs = []
        hand_obs.append(torch.cat([obs_buf[:, :self.num_hand_obs], obs_buf[:, 2*self.num_hand_obs:]], dim=1))
        hand_obs.append(torch.cat([obs_buf[:, self.num_hand_obs:2*self.num_hand_obs], obs_buf[:, 2*self.num_hand_obs:]], dim=1))
        rewards = self.task.rew_buf.unsqueeze(-1).to(self.rl_device)
        dones = self.task.reset_buf.to(self.rl_device)
        # Organize information into Multi-Agent RL format
        # Refer to https://github.com/tinyzqh/light_mappo/blob/HEAD/envs/env.py
        sub_agent_obs = []
        ...
        sub_agent_done = []
        for i in range(len(self.agent_index[0] + self.agent_index[1])):
            ...
            sub_agent_done.append(dones)
        # Transpose dim-0 and dim-1 values
        obs_all = torch.transpose(torch.stack(sub_agent_obs), 1, 0)
        ...
        done_all = torch.transpose(torch.stack(sub_agent_done), 1, 0)
        return obs_all, state_all, reward_all, done_all, info_all, None

    def reset(self):
        # Use a random action as the first action after the environment reset
        actions = 0.01 * (1 - 2 * torch.rand([self.task.num_envs, self.task.num_actions * 2], dtype=torch.float32, device=self.rl_device))
        # step the simulator
        self.task.step(actions)
        # Get the observation and state buffer in the environment, the detailed are the same as step(self, actions)
        obs_buf = torch.clamp(self.task.obs_buf, -self.clip_obs, self.clip_obs)
        ...
        obs = torch.transpose(torch.stack(sub_agent_obs), 1, 0)
        state_all = torch.transpose(torch.stack(agent_state), 1, 0)
        return obs, state_all, None

RL/Multi-Agent RL API

We also provide single-agent and multi-agent RL interfaces. In order to adapt to Isaac Gym and speed up the running efficiency, all operations are implemented on GPUs using tensor. Therefore, there is no need to transfer data between the CPU and GPU.

We give an example using HATRPO (the SOTA MARL algorithm for cooperative tasks) to illustrate multi-agent RL APIs, please refer to https://github.com/cyanrain7/TRPO-in-MARL:

from algorithms.marl.hatrpo_trainer import HATRPO as TrainAlgo
from algorithms.marl.hatrpo_policy import HATRPO_Policy as Policy
...
# warmup before the main loop starts
self.warmup()
# log data
start = time.time()
episodes = int(self.num_env_steps) // self.episode_length // self.n_rollout_threads
train_episode_rewards = torch.zeros(1, self.n_rollout_threads, device=self.device)
# main loop
for episode in range(episodes):
    if self.use_linear_lr_decay:
        self.trainer.policy.lr_decay(episode, episodes)
    done_episodes_rewards = []
    for step in range(self.episode_length):
        # Sample actions
        values, actions, action_log_probs, rnn_states, rnn_states_critic = self.collect(step)
        # Obser reward and next obs
        obs, share_obs, rewards, dones, infos, _ = self.envs.step(actions)
        dones_env = torch.all(dones, dim=1)
        reward_env = torch.mean(rewards, dim=1).flatten()
        train_episode_rewards += reward_env
        # Record reward at the end of each episode
        for t in range(self.n_rollout_threads):
            if dones_env[t]:
                done_episodes_rewards.append(train_episode_rewards[:, t].clone())
                train_episode_rewards[:, t] = 0

        data = obs, share_obs, rewards, dones, infos, \
                values, actions, action_log_probs, \
                rnn_states, rnn_states_critic
        # insert data into buffer
        self.insert(data)

    # compute return and update network
    self.compute()
    train_infos = self.train()
    # post process
    total_num_steps = (episode + 1) * self.episode_length * self.n_rollout_threads
    # save model
    if (episode % self.save_interval == 0 or episode == episodes - 1):
        self.save()

Testing

The trained model will be saved to logs/${Task Name}/${Algorithm Name}folder.

To load a trained model and only perform inference (no training), pass --test as an argument, and pass --model_dir to specify the trained models which you want to load. For single-agent reinforcement learning, you need to pass --model_dir to specify exactly what .pt model you want to load. An example of PPO algorithm is as follows:

python train.py --task=ShadowHandOver --algo=ppo --model_dir=logs/shadow_hand_over/ppo/ppo_seed0/model_5000.pt --test

For multi-agent reinforcement learning, pass --model_dir to specify the path to the folder where all your agent model files are saved. An example of HAPPO algorithm is as follows:

python train.py --task=ShadowHandOver --algo=happo --model_dir=logs/shadow_hand_over/happo/models_seed0 --test

Plotting

Users can convert all tfevent files into csv files and then try plotting the results. Note that you should verify env-num and env-step same as your experimental setting. For the details, please refer to the ./utils/logger/tools.py.

# geenrate csv for sarl and marl algorithms
$ python ./utils/logger/tools.py --alg-name <sarl algorithm> --alg-type sarl --env-num 2048 --env-step 8 --root-dir ./logs/shadow_hand_over --refresh 
$ python ./utils/logger/tools.py --alg-name <marl algorithm> --alg-type marl --env-num 2048 --env-step 8 --root-dir ./logs/shadow_hand_over --refresh 
# generate figures
$ python ./utils/logger/plotter.py --root-dir ./logs/shadow_hand_over --shaded-std --legend-pattern "\\w+"  --output-path=./logs/shadow_hand_over/figure.png

Use Bi-DexHands in Python scripts

import bidexhands as bi
import torch

env_name = 'ShadowHandOver'
algo = "ppo"
env = bi.make(env_name, algo)

obs = env.reset()
terminated = False

while not terminated:
    act = torch.tensor(env.action_space.sample()).repeat((env.num_envs, 1))
    obs, reward, done, info = env.step(act)

Enviroment Performance

Figures

We provide stable and reproducible baselins run by PPO, HAPPO, MAPPO, SAC algorithms. All baselines are run under the parameters of 2048 num_env and 100M total_step. The dataset folder contains the raw csv files.

ShadowHandOver ShadowHandLiftUnderarm
ShadowHandCatchUnderarm ShadowHandDoorOpenInward
ShadowHandOver2Underarm ShadowHandDoorOpenOutward
ShadowHandCatchAbreast ShadowHandDoorCloseInward
ShadowHandTwoCatchUnderarm ShadowHandDoorCloseOutward
ShadowHandPushBlock ShadowHandPen
ShadowHandScissors ShadowHandSwingCup
ShadowHandBlockStack ShadowHandReOrientation
ShadowHandPourWater ShadowHandSwitch
ShadowHandGraspAndPlace ShadowHandBottleCap

Offline RL Datasets

Data Collection

ppo_collect is the algo that collects offline data, which is basically the same as the mujoco data collection in d4rl. Firstly train the PPO for 5000 iterations, and collect and save the demonstration data in the first 2500 iterations:

python train.py --task=ShadowHandOver --algo=ppo_collection --num_envs=2048 --headless

Select model_5000.pt as the export policy to collect the expert dataset:

python3 train.py --task=ShadowHandOver --algo=ppo_collect --model_dir=./logs/shadow_hand_over/ppo_collect/ppo_collect_seed-1/model_5000.pt --test --num_envs=200 --headless

Similarly, select model.pt as the random policy, select a model as the medium policy, collect random data and medium data as above, and evenly sample the replay data set from the demonstration data before training to the medium policy. The size of each dataset is 10e6. Run merge.py to get the medium-expert dataset.

Offline Data

The originally collected data in our paper is available at: Shadow Hand Over, Shadow Hand Door Open Outward.

Use rl_games to train our tasks

Please note that we have only test to support rl-games==1.5.2. Higher or lower version may cause an error.

For example, if you want to train a policy for the ShadowHandOver task by the PPO algorithm, run this line in bidexhands folder:

python train_rlgames.py --task=ShadowHandOver --algo=ppo

Currently we only support PPO and PPO with LSTM methods in rl_games. If you want to use PPO with LSTM, run this line in bidexhands folder:

python train_rlgames.py --task=ShadowHandOver --algo=ppo_lstm

The log files using rl_games can be found in bidexhands/runs folder.

Known issue

It must be pointed out that Bi-DexHands is still under development, and there are some known issue:

Future Plan

Citation

Please cite as following if you think this work is helpful for you:

@inproceedings{
chen2022towards,
title={Towards Human-Level Bimanual Dexterous Manipulation with Reinforcement Learning},
author={Yuanpei Chen and Yaodong Yang and Tianhao Wu and Shengjie Wang and Xidong Feng and Jiechuan Jiang and Zongqing Lu and Stephen Marcus McAleer and Hao Dong and Song-Chun Zhu},
booktitle={Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track},
year={2022},
url={https://openreview.net/forum?id=D29JbExncTP}
}

The Team

Bi-DexHands is a project contributed by Yuanpei Chen, Yaodong Yang, Tianhao Wu, Shengjie Wang, Xidong Feng, Jiechuang Jiang, Hao Dong, Zongqing Lu, Song-chun Zhu at Peking University, please contact yaodong.yang@pku.edu.cn if you are interested to collaborate.

We also thank the list of contributors from the following two open source repositories: Isaac Gym, HATRPO.

We also recommend users to read the early work on dexterous hands manipulation that inpisres this work.

License

Bi-DexHands has an Apache license, as found in the LICENSE file.