ARISE-Initiative / robomimic

robomimic: A Modular Framework for Robot Learning from Demonstration
MIT License
592 stars 181 forks source link

How to train on bc_transformer? #162

Closed BUAAZhangHaonan closed 2 months ago

BUAAZhangHaonan commented 3 months ago

Thank you for your excellent work! I have a question about how to train on bc_transformer correctly? In script/train.py, we use run_epoch() to execute the training process, and we do this as follows:

# load next batch from data loader
        try:
            t = time.time()
            batch = next(data_loader_iter)
        except StopIteration:
            # reset for next dataset pass
            data_loader_iter = iter(data_loader)
            t = time.time()
            batch = next(data_loader_iter)
        timing_stats["Data_Loading"].append(time.time() - t)

        # process batch for training
        t = time.time()
        input_batch = model.process_batch_for_training(batch)
        input_batch = model.postprocess_batch_for_training(input_batch, obs_normalization_stats=obs_normalization_stats)
        timing_stats["Process_Batch"].append(time.time() - t)

In this way, we can get the data one by one from the dataloader through the next() operation, which means the shape of batch["actions"] is [_, 1, _]. In process_batch_for_training of BC_Transformer, batch["actions"] is used in this way:

if self.supervise_all_steps:
            # supervision on entire sequence (instead of just current timestep)
            input_batch["actions"] = batch["actions"][:, :h, :]
        else:
            # just use current timestep
            input_batch["actions"] = batch["actions"][:, h-1, :]

and the h is self.context_length which is not 1 and leads to a mistake as the default parm context_length in bc.json is 10:

IndexError: index 9 is out of bounds for dimension 1 with size 1

How can I adjust the code to fit the struct of bc_transformer?

amandlek commented 2 months ago

Please see our guide here - you can compare with the config there to see how to change the data loading options in the config.

BUAAZhangHaonan commented 2 months ago

My oversight, I didn't notice this tutorial. I tried it, it works! Thank you for this!