X-Diffusion

Public training code for X-Diffusion.

This first pass exposes classifier training on the same unified H5 data layout used by the policy-training code. The classifier script is kept close to the original H5 classifier path: it builds the H5 dataloaders, instantiates the diffusion policy modules, and trains HumanRobotClassifier.

Install

Create an environment with PyTorch, then install the Python dependencies:

pip install -r requirements.txt

Train Classifier

Edit the H5 data roots in configs/classifier_2class_alldata.yaml, then run:

python scripts/train_classifier.py --config_path configs/classifier_2class_alldata.yaml

For a small smoke run, start from configs/classifier_example.yaml and point retargeted_root, rgb_seg_root, object_data_root, and split_root at your local processed H5 data.

Outputs are written under:

classifier_checkpoints/<policy_and_image_mode>/<task>/<data_types>/<timestamp>/

Each run writes metrics.jsonl, config.yaml, dataset_stats/, checkpoints/latest.pth, and final noise-validation JSON when noise validation is enabled. WandB is disabled in the example configs; set use_wandb: true to enable it.

Data Layout

The classifier uses the same H5 roots as policy training:

<retargeted_root>/<task>/<robot|human_f|human_inf>/demo*.h5
<object_data_root>/<task>/<robot|human_f|human_inf>/demo*.h5
<rgb_seg_root>/<task>/<robot|human_f|human_inf>_imgs/demo*_imgs.h5

The default configs expect split files under split_root:

<split_root>/small_split_train_file_names.json
<split_root>/full_split_train_file_names.json

Physical H5 files should include keys such as ee_pos, ee_euler, gripper_open, and 3d_tracks. Object H5 files should include 3d_object_tracks. RGB/segmentation H5 files are part of the unified layout; image loading is disabled in the provided classifier configs with use_image: false.

Human data (human_f, human_inf) is labeled 0; robot data is labeled 1.