TORA: Topological Representation Alignment
for 3D Shape Assembly

Nahyuk Lee*^,1, Zhiang Chen*^,2, Marc Pollefeys², Sunghwan Hong^2,3

¹ Independent Researcher, ² ETH Zurich, ³ ETH AI Center   |   * denotes equal contribution

TL;DR: Assemble unposed 3D parts into complete objects using rectified flow with teacher-guided topological representation alignment.

Table of Contents

• Overview
• Getting Started
  • Setup
  • Teacher Weights
  • TORA Checkpoints
• Usage
  • Data Preparation
  • Training
  • Evaluation
• Advanced
  • Configurations
  • FAQ
• Citation
• Acknowledgments

Overview

TORA is a topology-first representation alignment framework that distills relational structure from a frozen pretrained 3D encoder into the flow-matching backbone during training. We then extend to employ a Centered Kernel Alignment (CKA) loss to match the similarity structure between student and teacher representations for enhanced topological alignment. Through systematic probing of diverse 3D encoders, we show that geometry- and contact-centric teacher properties, not semantic classification ability, govern alignment effectiveness, and that alignment is most beneficial at later transformer layers where spatial structure naturally emerges.

Key features

• Rectified Flow for point cloud assembly
• 6 teacher encoders: Uni3D (base/large/giant), Concerto, Find3D, OpenShape, PatchAlign3D, Sonata
• Projector: MLP
• 3 alignment losses: CKA, Cosine, NT-Xent
• Anchor-free and anchor-based assembly modes

Getting Started

Setup

Clone the repository with submodules:

git clone --recursive https://github.com/<org>/tora.git
cd tora

Create a conda environment with Python 3.10:

conda create -n tora python=3.10 -y
conda activate tora

Install dependencies:

bash install.sh

This environment includes PyTorch 2.5.1, PyTorch3D 0.7.8, and Flash Attention 2.7.4 with CUDA 12.4.

Teacher Weights

Model	Size	Description
uni3d_base.pt	178 MB	Uni3D base encoder
uni3d_large.pt	615 MB	Uni3D large encoder
uni3d_giant.pt	2.0 GB	Uni3D giant encoder
openshape_model.pt	388 MB	OpenShape encoder
patchalign3d_stage2.pt	89 MB	PatchAlign3D encoder

Teacher weights are stored in pretrained/ and downloaded automatically when needed.

TORA Checkpoints

Pretrained TORA checkpoints are provided in Hugging Face.

Usage

Data Preparation

TORA uses HDF5 format for all datasets. Each file should contain point clouds, normals, per-part counts, and anchor information.

Dataset	Source	Note
BBad-Everyday, PartNet, TwoByTwo	RPF	direct download
BBad-Artifact	BBad Official	needs conversion
Fantastic Breaks	Fantastic Breaks	needs conversion
Fractura	GARF	needs symlink setup

Conversion scripts are provided in dataset_process/:

# Convert Breaking Bad artifact subset
python dataset_process/convert_artifact_to_h5_igl.py

# Convert Fantastic Breaks
python dataset_process/convert_fantastic_breaks_to_h5.py

# Setup Fractura symlinks
bash dataset_process/setup_fractura.sh

Training

TORA follows the flow model training pipeline as in RPF (encoder pretraining + flow model training) with a frozen pretrained encoder, but adds representation alignment with frozen teacher models. We recommend training with representation alignment enabled, which is TORA's main contribution.

python train.py \
    data_root="../dataset" \
    data=main/bbad_everyday \
    model/teacher=uni3d \
    model/projector=mlp \
    model/alignment_loss=cka \
    trainer.devices=8

Tip

See RPF's training instructions for details on the base two-stage pipeline (encoder pretraining and flow model training). TORA uses the same pretrain.yaml config for encoder pretraining.

Available configuration options

Teachers: uni3d (default), concerto, find3d, openshape, patchalign3d, sonata

Projectors: mlp

Alignment Losses: cka (default), cosine, nt_xent

Datasets: main/bbad_everyday, main/partnet, main/twobytwo, main/bbad_artifact, main/fantastic_breaks, main/fractura

Override any parameter from the command line:

# Different teacher and loss
python train.py model/teacher=concerto model/alignment_loss=cosine

# Resume from checkpoint
python train.py ckpt_path=./output/TORA_base/last.ckpt

# Adjust learning rate and batch size
python train.py model.optimizer.lr=1e-4 data.batch_size=32

Evaluation

Run evaluation on a trained model:

python sample.py \
    ckpt_path=./output/TORA_base/best.ckpt \
    data_root="../dataset" \
    data=main/bbad_everyday

Metrics reported: Chamfer Distance, Part Accuracy, Rotation/Translation RMSE, Recall@5/10 (rotation), Recall@1cm/5cm (translation).

Advanced

Configurations

Hydra configuration structure

config/
├── train.yaml              # Training entry config
├── sample.yaml             # Inference/evaluation config
├── pretrain.yaml           # Encoder pretraining config
├── data/
│   ├── main/               # Main datasets (bbad_everyday, partnet, twobytwo, ...)
│   ├── rpf/                # RPF-compatible dataset configs
│   └── zeroshot/           # Zero-shot datasets (artifact, fantastic_breaks, fractura)
├── model/
│   ├── tora.yaml           # Main TORA model config
│   ├── feature_extractor/  # Encoder configs (ptv3_object)
│   ├── flow_model/         # DiT architecture config
│   ├── teacher/            # Teacher encoder configs (6 options)
│   ├── projector/          # Projector config (mlp)
│   └── alignment_loss/     # Alignment loss configs (3 options)
├── trainer/                # Lightning Trainer configs
└── loggers/                # Logger configs (wandb)

Key parameters in train.yaml:

• seed: Random seed (default: 42)
• data_root: Path to HDF5 dataset directory
• experiment_name: Name for WandB and output directory
• model.encoder_ckpt: Path to pretrained encoder checkpoint
• model.flow_model_ckpt: Path to flow model checkpoint for finetuning

FAQ

Note

FAQ will be updated. Please refer to GitHub Issues for questions in the meantime.

Citation

@article{lee2026tora,
  title     = {TORA: Topological Representation Alignment for 3D Shape Assembly},
  author    = {Lee, Nahyuk and Chen, Zhiang and Pollefeys, Marc and Hong, Sunghwan},
  journal   = {arXiv preprint arXiv:2604.04050},
  year      = {2026}
}

Acknowledgments

This codebase builds upon Rectified Point Flow and is inspired by REPA for representation alignment. We thank the authors of both works for their contributions.

License

See the LICENSE file for details.