real-stanford / diffusion_policy

[RSS 2023] Diffusion Policy Visuomotor Policy Learning via Action Diffusion
https://diffusion-policy.cs.columbia.edu/
MIT License
1.39k stars 260 forks source link

Possible Issue with Data Organization in Buffer Addition of kitchen_lowdim_dataset.py #61

Open hoost14 opened 5 months ago

hoost14 commented 5 months ago

Hello,

Thank you so much for your amazing work and beautiful code.

I'm currently working with Franka Kitchen dataset where the dimensions of a numpy array masks suggest that there are 566 demonstrations, each represented by a column in a (409, 566) shape array. Each column in masks corresponds to the existence of data per timestep for each demonstration.

However, I have encountered what seems to be a potential issue in the code where episodes are being added to the replay buffer. Here is the snippet:

https://github.com/real-stanford/diffusion_policy/blob/548a52bbb105518058e27bf34dcf90bf6f73681a/diffusion_policy/dataset/kitchen_lowdim_dataset.py#L23-L37

data_directory = pathlib.Path(dataset_dir) observations = np.load(data_directory / "observations_seq.npy") actions = np.load(data_directory / "actions_seq.npy") masks = np.load(data_directory / "existence_mask.npy")

self.replay_buffer = ReplayBuffer.create_empty_numpy() for i in range(len(masks)): eps_len = int(masks[i].sum()) obs = observations[i,:eps_len].astype(np.float32) action = actions[i,:eps_len].astype(np.float32) data = {
'obs': obs, 'action': action } self.replay_buffer.add_episode(data)

From this code, it appears that each iteration of the loop is supposed to handle a single demonstration. However, the indexing used (observations[i,:eps_len] and actions[i,:eps_len]) seems to imply that the demonstrations are organized by rows rather than columns. If each demonstration is indeed a column in the observations and actions arrays, the correct indexing should possibly be column-wise rather than row-wise.

Could you please confirm whether the demonstrations are intended to be represented as rows or columns in the dataset? If they are indeed columns, would the correct approach be to modify the indexing to reflect this structure?

Thank you for looking into this matter. I am looking forward to your clarification.

Best regards