Policy Training#

Run the following command to train WB-VIMA policies.

python3 main/train/train.py data_dir=<HDF5_PATH> \
    bs=<BS> \
    arch=wbvima \
    task=<TASK_NAME> \
    exp_root_dir=<EXP_ROOT_DIR> \
    gpus=<NUM_GPUS> \
    use_wandb=<USE_WANDB> \
    wandb_project=<WANDB_PROJECT>

We explain the arguments below.

  • data_dir: The path to the merged hdf5 file containing the training data.

  • bs: Batch size.

  • arch: The architecture to use. It must exist in the arch behavior-robot-suite/brs-algo folder. Now we use wbvima.

  • task: The task name. It must exist in the task behavior-robot-suite/brs-algo folder. Select one from clean_house_after_a_wild_party, clean_the_toilet, take_trash_outside, put_items_onto_shelves, and lay_clothes_out.

  • exp_root_dir: The directory to save the experiment results, including curves and model checkpoints.

  • gpus: The number of GPUs to use. Refer to the PyTorch Lightning documents for how to set it.

  • use_wandb: Whether to use Wandb for logging. Set it to True if you want to use it.

  • wandb_project: The Wandb project name. If Wandb is not used, set it to null.

For more parameters, refer to the training config file behavior-robot-suite/brs-algo. We follow the Hydra CLI syntax.