ARISE-Initiative / robomimic

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

Is there a way to get a full demonstration from start to finish. #168

Closed sarthakdas closed 2 months ago

sarthakdas commented 2 months ago

If I was to set seq_length to 100 would that guarantee a s0?

amandlek commented 2 months ago

We don't have out-of-the-box support for this, but you could probably achieve this by taking a few steps. Here's my recommendation:

  1. Modify the dataset hdf5 to pad all trajectories to a maximum length according to the maximum length trajectory in the dataset, and store the masks, so you know which timesteps are valid for each trajectory.
  2. In your robomimic config, disable dataset loader padding (e.g. pad_seq_length), and set the seq_length to the maximum trajectory length. Hopefully this should give you a batch of full trajectories, but you should test this out.
  3. Make sure the masks from step (1) are also loaded, so you only train over valid timesteps.

Hope this is helpful!