CEML

Cross-Embodiment Manipulation Learning (CEML) is a framework for training and evaluating robotic manipulation policies while varying the embodiment of the robot.

The project was developed as part of the Master's Thesis of Federico Morro, conducted at the Technical University of Munich (TUM) to then graduate with honors at Politecnico di Torino (PoliTO). The thesis was supervised by Prof. Giuseppe Bruno Averta (PoliTO) and Dr.rer.nat. Zhenshan Bing (TUM).

The complete thesis document is available here.

Repository Organization

The repository is structured as follows:

• ceml/: Implementation folder.
  • ceml/: Main Python package containing the implementation of the CEML state-based agents and the vision-based policy.
    • additional_modules/: Additional neural network modules (additional tokenizers and heads, weighted BC) for the vision-based policy.
    • agents/: AWAC-inspired BC + SAC RL state-based agents training and evaluation scripts.
    • datasets/: Dataset utilities (augmentation, HDF5 merging, dataset loader).
    • gripper_object/: Encoders for gripper and object information, used as input to the vision-based policy.
    • grippers/: Gripper point cloud extractiona and processing.
    • octo/: Implementation of the vision-based policy, based on the Octo generalist agent. The implementation is based on a PyTorch reimplementation of the original JAX code, available at octo-pytorch.
      • data/: Dataset loading and transforms.
        • utils/: Data utilities (goal relabeling, task augmentation, text processing).
      • model/: PyTorch implementation of the Octo model.
        • components/: Architecture components (action heads, block-attention transformer, tokenizers, ViT encoders).
      • utils/: Utilities for training, gym wrappers, and RLDS dataset handling.
    • offline_rl/: Algorithms for offline reinforcement learning (AWR, ReBRAC).
    • online_finetune/: Algorithms for online fine-tuning (DPPO, IBRL).
    • representation_learning/: Tools for latent space structuring (constrastive additional losses, latent space visualization, demo comparison).
    • utils/: General utilities (seed management, additional dataset utils, robosuite utils).
    • visualize/: Visualization tools (agents performance, dataset metadata).
    • wrappers/: Environment wrappers for Robosuite and GG-CNN integration.
    • octo_jax/: Original JAX scripts.
    • out/: Output folder for training logs, evaluation results, and visualizations.
    • scripts/: Scripts for training, evaluation, and data processing of the vision-based policy.
      • configs/: YAML configuration files for training and evaluation runs.
    • tests/: Tests for some components of the codebase.
• datasets/: Scripts for downloading and preprocessing openly available datasets (e.g., Cornell Grasping, Robomimic).
• ggcnn/: Derived from the GG-CNN repository (PyTorch port), used for grasp detection.
• robomimic/: Derived from Robomimic, providing the datasets and utilities for offline RL training.
• robosuite/: Derived from Robosuite, providing the simulation environment.
• run_scripts/: Specialized bash scripts for running training and evaluation jobs (Slurm/local).
• weights/: Scripts to download pre-trained model weights.
• utils/: Miscellaneous utility scripts.

Useful Commands

Setup Container

The code was run inside a custom container on both local machine and on a remote AI cluster (LRZ) using Enroot and Slurm. The following commands refer to both environments, and might be adapted accordingly.

• Create container, extending NVIDIA NGC base image with Ubuntu 20.04, and mounting the CEML repository inside /home

enroot create --name ceml_cnt nvidia+base+ubuntu+20.04_x64_2022-09-23.sqsh
enroot start --root --mount CEML:/home ceml_cnt

• Inside the container, run container_setup.sh to install dependencies and create the ceml_env Conda environment

cd /home
./container_setup.sh

• Exit the container and export the modified container as a new SquashFS image

exit
enroot export --output ceml_img.sqsh ceml_cnt

Re-Use Container

enroot create --name ceml_cnt ceml_img.sqsh
enroot start --root --mount CEML:/home ceml_cnt
conda activate ceml_env

Troubleshooting

Error in DLiMP Dataset Loading

ValueError: Error while flattening dict: FeaturesDict received a non dict item: <tensorflow_datasets.core.decode.base.SkipDecoding object at 0x7f2ca83aeda0>

vim /opt/conda/envs/ceml_env/lib/python3.10/site-packages/dlimp/dataset.py
... comment out line 147 ...
        #decoders={"steps": tfds.decode.SkipDecoding()},