# 3D-aware Disentangled Representation for Compositional Reinforcement Learning (Anonymized)

This repository contains the anonymized code used to train agents for object manipulation in Isaac Gym. Personal identifiers and publication links have been removed for double-blind review while keeping the full training and data-processing pipeline intact. It builds on the ECRL repository (MIT) and incorporates code from the OSRT repository (MIT) under the same license.

## Requirements

| Library             | Version (tested) |
|---------------------|------------------|
| Python              | 3.8              |
| torch               | 2.1.2            |
| stable-baselines3   | 1.5.0            |
| isaacgym            | Preview build    |

Install Isaac Gym from the official preview release and follow its installation instructions. Additional Python dependencies are listed in `requirements.txt`.

## Training
### Pre-training 3D block-slot representation

```bash
python -m torch.distributed.launch \
    --use_env \
    --standalone \
    --nnodes 1 \
    --nproc_per_node 2 \
    train_bsa.py \
    <path_to_config_dir>
```
This script trains the 3D block-slot representation model. It expects the same config directory structure used for RL (contains `Config.yaml`, `IsaacPandaPushConfig.yaml`, and `PolicyConfig.yaml`); pass it via `<path_to_config_dir>`.

### Goal-conditioned Reinforcement Learning
The code supports table-top pushing tasks. Each experiment expects a configuration directory with three files: `Config.yaml`, `IsaacPandaPushConfig.yaml`, and `PolicyConfig.yaml`. Point `--config_dir` to the desired directory when running training.

```bash
python main.py -c <path_to_config_dir>
```

- `Config.yaml` controls agent hyperparameters and logging.
- `IsaacPandaPushConfig.yaml` controls simulator parameters (number of objects, camera setup, table sizes, etc.).
- `PolicyConfig.yaml` provides policy-network definitions for the chosen method.

## Repository layout
- `dlp2/`: deep latent particles representation code and utilities.
- `bsa/`: 3D block-slot representation and OSRT code and utilities (adapted from OSRT).
- `pre_trained_models/`: placeholder directory for pretrained representation checkpoints.

- `main.py`: entry point for RL training and data collection.
- `td3_agent.py`, `policies.py`, `multi_her_replay_buffer.py`: TD3+HER implementation and policy variants (Block transformer, EIT).
- `isaac_panda_push_env.py`, `isaac_vec_env.py`, `isaac_env_wrappers.py`, `panda_controller/`: Isaac Gym environment definitions and controllers.
- `assets/urdf`: robot and object assets for simulation.

## Third-party code and licenses

This project includes code and assets adapted from external sources. See `THIRD_PARTY_LICENSES.md` for a summary and `assets/licenses/` for bundled license files. Original license headers are retained in the source files where required. STL object meshes under `assets/object/` originate from the CATER dataset (Apache-2.0).

## Citation
Citation details will be added after the review process.
