bugs about data_gen.py

[LemonWade]

    def __getitem__(self, index: int) -> None:
        ...
        try:
            demo, state_ls, action_ls = get_observation(
                task, variation, episode, self.env
            )
        except (FileNotFoundError, RuntimeError, IndexError, EOFError) as e:
            print(e)
            return

        state_ls = einops.rearrange(
            state_ls,
            "t 1 (m n ch) h w -> t n m ch h w",
            ch=3,
            n=len(args.cameras),
            m=2,
        )

        frame_ids = list(range(len(state_ls) - 1))
        num_frames = len(frame_ids)
        attn_indices = get_attn_indices_from_demo(task, demo, args.cameras)

        if (task in self.variable_lengths and num_frames > self.max_eps_dict[task]) or (
            task not in self.variable_lengths and num_frames != self.max_eps_dict[task]
        ):
            print(f"ERROR ({task}, {variation}, {episode})")
            print(f"\t {len(frame_ids)} != {self.max_eps_dict[task]}")
            return

        state_dict: List = [[] for _ in range(5)]
        print("Demo {}".format(episode))
        state_dict[0].extend(frame_ids)
        state_dict[1].extend(state_ls[:-1])
        state_dict[2].extend(action_ls[1:])
        state_dict[3].extend(attn_indices)
        state_dict[4].extend(action_ls[:-1])  # gripper pos

        np.save(taskvar_dir / f"ep{episode}.npy", state_dict)  # type: ignore

I encountered an issue in the __getitem__method of a dataset processing script. The line frames.insert(0, 0) adds an element at position 0, but after attn_indices = get_attn_indices_from_demo(task, demo, args.cameras), there's no further manipulation of attn_indices. This discrepancy leads to an error when saving the data using np.save(taskvar_dir / f"ep{episode}.npy", state_dict). The error message is as follows:

ValueError: setting an array element with a sequence. The requested array has an inhomogeneous shape after 1 dimensions. The detected shape was (5,) + inhomogeneous part.

How can I modify the script to resolve this issue?

The problem seems to arise from inconsistent data shapes within state_dict, especially regarding attn_indices, which does not align with the modifications made to other data sequences in the script. Any suggestions or corrections would be greatly appreciated.

[rakhimovv]

I managed to fix it by replacing this line with np.save(taskvar_dir / f"ep{episode}.npy", np.array(state_dict, dtype=object), allow_pickle=True)

[LemonWade]

"Following your advice, I have successfully run it. Thank you very much."

[nickgkan]

For everyone coming back here from the future, please give a shot to our updated data generation scripts.