Official PyTorch code release of the paper "Entity-Centric Reinforcement Learning for Object Manipulation from Pixels" by Dan Haramati, Tal Daniel and Aviv Tamar.

Abstract

Manipulating objects is a hallmark of human intelligence, and an important task in domains such as robotics. In principle, Reinforcement Learning (RL) offers a general approach to learn object manipulation. In practice, however, domains with more than a few objects are difficult for RL agents due to the curse of dimensionality, especially when learning from raw image observations. In this work we propose a structured approach for visual RL that is suitable for representing multiple objects and their interaction, and use it to learn goal-conditioned manipulation of several objects. Key to our method is the ability to handle goals with dependencies between the objects (e.g., moving objects in a certain order). We further relate our architecture to the generalization capability of the trained agent, based on a theoretical result for compositional generalization, and demonstrate agents that learn with 3 objects but generalize to similar tasks with over 10 objects.

Haramati, Dan, Tal Daniel, and Aviv Tamar. "Entity-Centric Reinforcement Learning for Object Manipulation from Pixels." Proceedings of the Twelfth International Conference on Learning Representations (ICLR). 2024.

 @inproceedings{
 haramati2024entitycentric,
 title={Entity-Centric Reinforcement Learning for Object Manipulation from Pixels},
 author={Dan Haramati and Tal Daniel and Aviv Tamar},
 booktitle={The Twelfth International Conference on Learning Representations},
 year={2024},
 url={https://openreview.net/forum?id=uDxeSZ1wdI}
 }

Entity Centric-Reinforcement Learning (ECRL)

1. Prerequisites
2. Environments
3. Training
   • Deep Latent Particles (DLP) Pretraining
   • RL Training
4. Evaluation
5. Repository Content Details
6. Credits

The following are the main libraries required to run this code:

For the full list of requirements, see the requirements.txt file.

For the simulation environment, download the Isaac Gym Preview release from the website, then follow the installation instructions in the documentation.

The above figure describes the suite of training environments used in the paper.

N-Cubes: Push N different-colored cubes to their goal location.
Adjacent-Goals: A 3-Cubes setting where goals are sampled randomly on the table such that all cubes are adjacent. This task requires accounting for interactions between objects.
Ordered-Push: A 2-Cubes setting where a narrow corridor is set on top of the table such that its width can only fit a single cube. We consider two possible goal configurations: red cube in the rear of the corridor and green cube in the front, or vice versa. This task requires to fulfill the goals in a certain order, otherwise the agent fails (pulling a block out of the corridor is not possible).
Small-Table: A 3-Cubes setting where the table is substantially smaller. This task requires to accurately account for all objects in the scene at all times, to avoid pushing blocks off the table.
Push-2T: Push 2 T-shaped blocks to a single goal orientation.

A configuration file for each environment, IsaacPandaPushConfig.yaml, is provided in the corresponding directory in config.

We provide pretrained model checkpoints:

Download and place in the relevant directory in latent_rep_chkpts (e.g., checkpoint of DLP trained on data from the 5-Cubes environment should be placed in latent_rep_chkpts/dlp_push_5C).

In order to retrain the model:

1. Collect image data using a random policy by running main.py -c <configuration_dir> with the desired environment (e.g, main.py -c config/n_cubes), setting collectData: True and collectDataNumTimesteps in the relevant Config.yaml. This will save a .npy file in the results directory.
2. Process the data into a dataset by running dlp2/datasets/process_dlp_data.py (fill in the relevant paths in the beginning of the script).
3. Configure config/TrainDLPConfig.yaml and run train_dlp.py.

Run main.py -c <configuration_dir> with the desired configuration (e.g, main.py -c config/n_cubes).
Config.yaml contains agent and training parameters and IsaacPandaPushConfig.yaml contains environment parameters.

In order to reproduce the experiments in the paper, input the corresponding configuration directory. The configurations are already set to match the ones used in the paper. The parameters requiring configuration for the different instances of the experiments (e.g, 'State' or 'Image'):

• In Config.yaml the Model parameters.
• In IsaacPandaPushConfig.yaml the numObjects parameter (for n_cubes and push_t).

To log training statistics and images/videos using Weights & Biases set WANDB: log: True in Config.yaml and fill in your username in the wandb.init(entity="") line in the main.py script.

Agent model checkpoints and intermediate results are saved in the model_chkpts and results directories respectively.

To evaluate an agent on a given environment, run policy_eval.py. Set the agent model_path and the desired configuration directory manually in the beginning of the script.

Evaluation on Zero-Shot Generalization

Cube Sorting: train on config/n_cubes with numObjects: 3 and evaluate on config/generalization_sort_push.

Different Number of Cubes than in Training: train on config/generalization_num_cubes with numObjects: 3 and evaluate with same config and varying number of objects.

isaac_env_wrappers.py, isaac_panda_push_env.py, and chamfer_reward.py contain scripts for debugging the components seperately.
latent_classifier.py contains an interactive script for tagging data and training the classifier.

• Environments adapted from the IsaacGymEnvs FrankaCubeStack environment.
• Controller for the Franka Panda arm adapted from the OSC implementation of OSCAR.
• RL code modified from SB3 (mainly from the TD3 agent and HERReplayBuffer).
• Entity Interaction Transformer (EIT) code built from components adapted from the DDLP Particle Interaction Transformer (PINT) which is based on minGPT.
• SMORL re-implementation based on the official repository of the paper using the same codebase as the EIT.
• DLP code modified from the official implementation of DLPv2 in the DDLP repository.
• VAE code modified from Taming Transformers.
• Slot-Attention code modified from Object Discovery PyTorch.