POCR: Composing Pre-Trained Object-Centric Representations for Robotics From "What" and "Where" Foundation Models

International Conference on Robotics and Automation (ICRA), 2024

Junyao Shi*, Jianing Qian*, Jason Ma, Dinesh Jayaraman

University of Pennsylvania

This is the official repository for POCR, Pre-Trained Object-Centric Representations for Robotics. POCR chains "what" and "where" foundation models for robotic control. The “where” foundation model produces a set of segmentation masks representing object candidates in the scene. Slot binding selects which among them to bind to the slots in our object-centric representation. Image contents in each slot are represented by the “what” foundation model and their mask bounding box coordinates. The robot learns policies over slot representations.

This repository contains code that for how POCR generates and processes masks for robotic manipulation tasks using SAM (Segment Anything Model) and tracks them within and across videos using XMem.

Installation

1. Clone the Repository

git clone https://github.com/junyaoshi/POCR.git
cd POCR

2. Install XMem

XMem needs to be installed separately and referenced in the configuration. Please refer to their Github repo: https://github.com/hkchengrex/XMem. Clone the code and make a conda environment for running XMem.

3. Download SAM Model

Download the Segment Anything Model (SAM) checkpoint from here
Store the model file in your desired location
Update the model path in conf/method/pocr.yaml:

sam:
  checkpoint: "/path/to/sam_vit_h_4b8939.pth"

4. Set Up POCR Environment

# Create POCR conda environment
conda create -n pocr python=3.8
conda activate pocr
# Install requirements
pip install -r requirements.txt

Dataset Structure

Your dataset should follow this structure:

task_name/
├── 0/
│   ├── 0.png
│   ├── 1.png
│   ├── 2.png
│   └── ...
├── 1/
│   ├── 0.png
│   ├── 1.png
│   └── ...
└── ...

Each numbered directory represents an episode, and contains sequentially numbered PNG frames starting from 0.png.

Configuration

Update the paths in conf/generate_masks.yaml:

paths:
  data_root: "/path/to/your/dataset"
  debug_root: "/path/to/debug/output"
  xmem_root: "/path/to/your/XMem"
  queue_root: "/path/to/queue/output"
  log_root: "/path/to/logs"

Configure task-specific parameters in conf/task/your_task.yaml and specify the task in conf/generate_masks.yaml:

defaults:
  - _self_
  - task: task_name
  - method: pocr

Usage

Activate the POCR environment:

conda activate pocr

Run the mask generation:

python generate_masks.py

Replace task_name with your specific task (e.g., close_box, pick_up_cup, etc.).

Demo Dataset

A demo dataset is available here to help you get started.

Output Structure

The program will generate:

Mask predictions in the Annotations directory
Original frames in the JPEGImages directory
Debug visualizations if enabled in the configuration

Visualization

To visualize XMem results:

python tools/visualize_xmem.py --input_dir /path/to/results --output_dir /path/to/visualization

Input directory is the data directory that contains the folder named "Annotations".

Parameter Configuration Guide

Each task's configuration can be customized in conf/task/task_name.yaml. Below is a detailed explanation of the parameters:

SAM (Segment Anything Model) Parameters

pred_iou_thresh: 0.90    # Confidence threshold for mask predictions (0-1)
                        # Higher values produce more confident but fewer masks
stability_thres: 0.95   # Threshold for mask stability score (0-1)
                        # Higher values ensure more stable and consistent masks
nms_thres: 0.65        # Non-maximum suppression threshold (0-1)
                        # Controls overlap between masks, higher values allow less overlap

Background Region Detection

n_bg_episodes: 5       # Number of episodes to use for background analysis
n_bg_imgs: 20         # Number of images to sample for background detection
sample_interval: 5     # Interval for sampling points in background analysis
n_kmeans_clusters: 15  # Number of clusters for background segmentation
votes_thres: 0.0003   # Threshold for background region voting
                      # Lower values detect more potential background regions
votes_percentage: 30   # Percentage of votes needed to classify as background

Mask Filtering Parameters

# Area Constraints
area_lower_bound: 50    # Minimum pixel area for valid masks
area_upper_bound: 10000 # Maximum pixel area for valid masks
# Overlap Filtering
filter_curr_mask: True  # Enable filtering based on current mask overlap
curr_mask_thres: 0.8    # Maximum allowed overlap with current mask (0-1)
# Background Filtering
filter_bg: True         # Enable background filtering
bg_thres: 0.8          # Background overlap threshold (0-1)
# Robot Mask Handling
filter_robot: True      # Enable robot mask filtering
robot_thres: 0.5       # Robot overlap threshold (0-1)
robot_mask_id: 16      # ID for robot mask in output
robot_color: 125       # Color value for robot visualization
# General Settings
num_masks: 5           # Maximum number of masks to generate per frame

Parameter Tuning Guidelines

Mask Quality Control
- Increase pred_iou_thresh for higher quality but fewer masks
- Decrease stability_thres if masks are too fragmented
- Adjust nms_thres to control mask overlap
Background Detection
- Adjust votes_thres and votes_percentage based on scene complexity
- Increase n_bg_episodes and n_bg_imgs for more robust background detection
Size Filtering
- Modify area_lower_bound and area_upper_bound based on objects in your scene
- Typical objects should fall well within these bounds
Overlap Control
- Adjust curr_mask_thres to control mask separation
- Modify bg_thres and robot_thres based on your scene setup

Example Configuration

# Configuration optimized for close_box task
pred_iou_thresh: 0.90
stability_thres: 0.95
nms_thres: 0.65
area_lower_bound: 50
area_upper_bound: 10000
curr_mask_thres: 0.8
bg_thres: 0.8
robot_thres: 0.5
num_masks: 5

Common Issues

XMem Environment Issues
- Ensure XMem environment has all dependencies installed
- Check CUDA compatibility
Path Configuration
- Double-check all paths in generate_masks.yaml
- Ensure dataset follows the required structure

Contributing

Feel free to open issues or submit pull requests for any improvements.

Acknowledgments

XMem for video object segmentation
Segment Anything for mask generation