Clarification on RoboMimic suboptimal classifier training: mixed-task or per-task?

Start with curation/scripts/subop_train_robomimic.sh and the HDF5Dataset implementation described in the issue. Compare the paper's RoboMimic experiment setup with how data_dir recursively discovers and interleaves HDF5 files. Done means documenting whether classifiers are trained jointly or per task and showing the correct directory layout.

Hi, thanks for releasing the SCIZOR code.

I have a question about the RoboMimic suboptimal classifier training script:

curation/scripts/subop_train_robomimic.sh

In this script, the training data is specified by:

--config.hdf5_dataset_kwargs.data_dir PATH_TO_ROBOMIMIC_HDF5_DATASET_FOLDER

From my reading of the code, HDF5Dataset recursively walks through data_dir and loads all .hdf5 files under that directory:

for root, dirs, files in os.walk(self.hdf5_config['data_dir']):
    for file in files:
        if file.endswith('.hdf5'):
            self.paths.append(os.path.join(root, file))

Then all discovered HDF5 files are wrapped as sub-datasets and interleaved for training.

I would like to clarify the intended RoboMimic setup used in the paper:

For the RoboMimic experiments, was the suboptimal classifier trained jointly on multiple tasks, e.g. Can MH + Square MH, or was a separate suboptimal classifier trained for each task?

Should users point data_dir to a root folder containing multiple RoboMimic tasks, such as:

robomimic_data/
├── can/mh/*.hdf5
└── square/mh/*.hdf5

or should they run subop_train_robomimic.sh separately with task-specific folders, such as:

robomimic_data/can/mh/
robomimic_data/square/mh/

Thanks! This clarification would be very helpful for reproducing the RoboMimic results and for understanding how to set up data_dir correctly.