# ankile/robust-rearrangement architecture

> Architecture diagram of the GitHub repository ankile/robust-rearrangement, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-09-28). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/ankile/robust-rearrangement
- Repository: https://github.com/ankile/robust-rearrangement

## Overview

Robust Rearrangement is a robotics research codebase for learning and refining policies that assemble furniture, from demonstrations through simulation evaluation and real-robot use. Its main workflow collects demonstrations, converts them into training datasets, trains behavior-cloning or reinforcement-learning policies, then evaluates policies in simulated rollouts. A separate sim-to-real tooling branch rerenders demonstrations with randomized scenes. The graph reflects the README and sampled source; some subsystem links are documented at a high level, while unsampled implementation details are intentionally left unwired.

Main entry points: [`src/behavior/base.py`](https://github.com/ankile/robust-rearrangement/blob/main/src/behavior/base.py) (Behavior policies), [`src/dataset/dataset.py`](https://github.com/ankile/robust-rearrangement/blob/main/src/dataset/dataset.py) (Sequence datasets), [`src/train/residual_ppo_w_bc.py`](https://github.com/ankile/robust-rearrangement/blob/main/src/train/residual_ppo_w_bc.py) (PPO refinement), [`src/eval/rollout.py`](https://github.com/ankile/robust-rearrangement/blob/main/src/eval/rollout.py) (Policy rollout)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_collect["Demonstration data"]
  node_raw[("Raw demonstrations")]
  node_collector["Demo collection"]
  node_processing["Demo processing<br/>[process_pickles.py]"]
  node_zarr[("Processed demonstrations<br/>[zarr.py]")]
end

subgraph group_learn["Policy learning"]
  node_dataset["Sequence datasets<br/>[dataset.py]"]
  node_normalizer["Data normalization<br/>[normalizer.py]"]
  node_behavior["Behavior policies<br/>[base.py]"]
  node_models["Neural models<br/>[transformer.py]"]
  node_bc["BC training<br/>[bc.py]"]
  node_dagger["DAgger training<br/>[dagger.py]"]
  node_ppo["PPO refinement"]
  node_buffer[("Rollout buffer<br/>[rollout_buffer.py]")]
end

subgraph group_sim["Simulation evaluation"]
  node_simenv["Furniture simulator"]
  node_rollout["Policy rollout<br/>[rollout.py]"]
  node_evaluate["Evaluation runner<br/>[evaluate_model.py]"]
end

subgraph group_robot["Robot execution"]
  node_realpolicy["Real inference<br/>[minimal.py]"]
  node_teleop["Real teleoperation<br/>[teleop_sm.py]"]
end

subgraph group_transfer["Simulation transfer"]
  node_rerender["Demo rerendering"]
  node_assets["Part conversion"]
end

node_operator(("Operator"))
node_robot(("Physical robot"))

node_operator -->|"provides demonstrations"| node_collector
node_collector -->|"saves trajectories"| node_raw
node_raw -->|"loads pickles"| node_processing
node_processing -->|"writes datasets"| node_zarr
node_zarr -->|"loads samples"| node_dataset
node_dataset -->|"feeds training"| node_bc
node_dataset -->|"feeds training"| node_dagger
node_dataset -->|"loads demos"| node_ppo
node_bc -->|"trains policy"| node_behavior
node_dagger -->|"trains policy"| node_behavior
node_ppo -->|"refines policy"| node_behavior
node_behavior -->|"uses networks"| node_models
node_evaluate -->|"creates environment"| node_simenv
node_evaluate -->|"runs evaluation"| node_rollout
node_rollout -->|"queries policy"| node_behavior
node_rollout -->|"steps environment"| node_simenv
node_ppo -->|"collects experience"| node_simenv
node_ppo -->|"stores rollouts"| node_buffer
node_operator -->|"provides commands"| node_teleop
node_teleop -->|"controls robot"| node_robot
node_robot -->|"provides observations"| node_realpolicy
node_realpolicy -->|"sends actions"| node_robot
node_raw -.->|"reads demonstrations"| node_rerender
node_rerender -.->|"produces episodes"| node_zarr
node_assets -.->|"supplies geometry"| node_rerender

click node_collector "https://github.com/ankile/robust-rearrangement/blob/main/src/data_collection/data_collector_sm.py"
click node_processing "https://github.com/ankile/robust-rearrangement/blob/main/src/data_processing/process_pickles.py"
click node_zarr "https://github.com/ankile/robust-rearrangement/blob/main/src/dataset/zarr.py"
click node_dataset "https://github.com/ankile/robust-rearrangement/blob/main/src/dataset/dataset.py"
click node_normalizer "https://github.com/ankile/robust-rearrangement/blob/main/src/dataset/normalizer.py"
click node_behavior "https://github.com/ankile/robust-rearrangement/blob/main/src/behavior/base.py"
click node_models "https://github.com/ankile/robust-rearrangement/blob/main/src/models/transformer.py"
click node_bc "https://github.com/ankile/robust-rearrangement/blob/main/src/train/bc.py"
click node_dagger "https://github.com/ankile/robust-rearrangement/blob/main/src/train/dagger.py"
click node_ppo "https://github.com/ankile/robust-rearrangement/blob/main/src/train/residual_ppo_w_bc.py"
click node_buffer "https://github.com/ankile/robust-rearrangement/blob/main/src/dataset/rollout_buffer.py"
click node_simenv "https://github.com/ankile/robust-rearrangement/blob/main/src/gym/mj_dual_franka_env.py"
click node_rollout "https://github.com/ankile/robust-rearrangement/blob/main/src/eval/rollout.py"
click node_evaluate "https://github.com/ankile/robust-rearrangement/blob/main/src/eval/evaluate_model.py"
click node_realpolicy "https://github.com/ankile/robust-rearrangement/blob/main/src/real/minimal.py"
click node_teleop "https://github.com/ankile/robust-rearrangement/blob/main/src/real/teleop_sm.py"
click node_rerender "https://github.com/ankile/robust-rearrangement/blob/main/src/sim2real/isaac_lab_rerender.py"
click node_assets "https://github.com/ankile/robust-rearrangement/blob/main/src/sim2real/convert_parts_to_usd.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_raw,node_collector,node_processing,node_zarr toneBlue
class node_dataset,node_normalizer,node_behavior,node_models,node_bc,node_dagger,node_ppo,node_buffer toneAmber
class node_simenv,node_rollout,node_evaluate toneMint
class node_realpolicy,node_teleop toneRose
class node_rerender,node_assets,node_operator,node_robot toneIndigo
```

## Components

### Demonstration data

- **Raw demonstrations**
- [Demo collection](https://github.com/ankile/robust-rearrangement/blob/main/src/data_collection/data_collector_sm.py): `src/data_collection/data_collector_sm.py`
- [Demo processing](https://github.com/ankile/robust-rearrangement/blob/main/src/data_processing/process_pickles.py): `src/data_processing/process_pickles.py`
- [Processed demonstrations](https://github.com/ankile/robust-rearrangement/blob/main/src/dataset/zarr.py): `src/dataset/zarr.py`

### Policy learning

- [Sequence datasets](https://github.com/ankile/robust-rearrangement/blob/main/src/dataset/dataset.py): `src/dataset/dataset.py`
- [Data normalization](https://github.com/ankile/robust-rearrangement/blob/main/src/dataset/normalizer.py): `src/dataset/normalizer.py`
- [Behavior policies](https://github.com/ankile/robust-rearrangement/blob/main/src/behavior/base.py): `src/behavior/base.py`
- [Neural models](https://github.com/ankile/robust-rearrangement/blob/main/src/models/transformer.py): `src/models/transformer.py`
- [BC training](https://github.com/ankile/robust-rearrangement/blob/main/src/train/bc.py): `src/train/bc.py`
- [DAgger training](https://github.com/ankile/robust-rearrangement/blob/main/src/train/dagger.py): `src/train/dagger.py`
- [PPO refinement](https://github.com/ankile/robust-rearrangement/blob/main/src/train/residual_ppo_w_bc.py): `src/train/residual_ppo_w_bc.py`
- [Rollout buffer](https://github.com/ankile/robust-rearrangement/blob/main/src/dataset/rollout_buffer.py): `src/dataset/rollout_buffer.py`

### Simulation evaluation

- [Furniture simulator](https://github.com/ankile/robust-rearrangement/blob/main/src/gym/mj_dual_franka_env.py): `src/gym/mj_dual_franka_env.py`
- [Policy rollout](https://github.com/ankile/robust-rearrangement/blob/main/src/eval/rollout.py): `src/eval/rollout.py`
- [Evaluation runner](https://github.com/ankile/robust-rearrangement/blob/main/src/eval/evaluate_model.py): `src/eval/evaluate_model.py`

### Robot execution

- [Real inference](https://github.com/ankile/robust-rearrangement/blob/main/src/real/minimal.py): `src/real/minimal.py`
- [Real teleoperation](https://github.com/ankile/robust-rearrangement/blob/main/src/real/teleop_sm.py): `src/real/teleop_sm.py`

### Simulation transfer

- [Demo rerendering](https://github.com/ankile/robust-rearrangement/blob/main/src/sim2real/isaac_lab_rerender.py): `src/sim2real/isaac_lab_rerender.py`
- [Part conversion](https://github.com/ankile/robust-rearrangement/blob/main/src/sim2real/convert_parts_to_usd.py): `src/sim2real/convert_parts_to_usd.py`

### Other components

- **Operator**
- **Physical robot**

## Connections

- Operator → Demo collection: provides demonstrations
- Demo collection → Raw demonstrations: saves trajectories
- Raw demonstrations → Demo processing: loads pickles
- Demo processing → Processed demonstrations: writes datasets
- Processed demonstrations → Sequence datasets: loads samples
- Sequence datasets → BC training: feeds training
- Sequence datasets → DAgger training: feeds training
- Sequence datasets → PPO refinement: loads demos
- BC training → Behavior policies: trains policy
- DAgger training → Behavior policies: trains policy
- PPO refinement → Behavior policies: refines policy
- Behavior policies → Neural models: uses networks
- Evaluation runner → Furniture simulator: creates environment
- Evaluation runner → Policy rollout: runs evaluation
- Policy rollout → Behavior policies: queries policy
- Policy rollout → Furniture simulator: steps environment
- PPO refinement → Furniture simulator: collects experience
- PPO refinement → Rollout buffer: stores rollouts
- Operator → Real teleoperation: provides commands
- Real teleoperation → Physical robot: controls robot
- Physical robot → Real inference: provides observations
- Real inference → Physical robot: sends actions
- Raw demonstrations → Demo rerendering: reads demonstrations
- Demo rerendering → Processed demonstrations: produces episodes
- Part conversion → Demo rerendering: supplies geometry