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

[j-cyoung]

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.