Updating and debugging sequence_dataset()

[HesNobi]

Hello,
Thank you very much for this open project and resources. I am researching Imitation and inverse reinforcement learning. D4RL has been a lifesaver.

The following code returns an error since some keys in the dataset arent np.ndarray or they don't present each step. I have fixed that by selecting the proper keys. Also to speed up the generator, I have added some limits on either the number of steps or episodes.

for k in dataset:
            data_[k].append(dataset[k][i])

I hope this pull request will be accepted
Thanks,

[jren03]

I seem to be erroring out on this line when I call d4rl.sequence_dataset(env) where env is any halfcheetah-medium-v2, hopper-medium-v2, or walker2d-medium-v2.

To reproduce:

d_iter = d4rl.sequence_dataset(env)
for traj in d_iter:
     pass

File "foo.py", line 111, in get_dataset
    for traj in d_iter:
File "/home/user/.conda/envs/d4rl/lib/python3.7/site-packages/d4rl/__init__.py", line 175, in sequence_dataset
    data_[k].append(dataset[k][i])
IndexError: invalid index to scalar variable.