How do you load proprio when using dataset mix?

[IrvingF7]

Hi team!

Have a quick question regarding the oxe dataloader

So, many OXE datasets have their own way of handling proprio data. Some use euler, and some use quaternion, etc. This would result in vector size differences, no? (7 for euler and 8 for quaternion)

Now, if I turn on load_proprio=True in make_oxe_dataset_kwargs_and_weights, I notice that when I load the datasets, I get many error messages like this:

W tensorflow/core/kernels/data/experimental/ignore_errors_dataset_op.cc:119] Error raised with error message Cannot add tensor to the batch: number of elements does not match. Shapes are: [tensor]: [2,7], [batch]: [2,8]

I understand that Octo and openvla are both trained without the proprio, so not being able to do this is fine. But what if I want to load the proprio data in a mix? Even just padding euler with 0 is fine.

I tried to add something in oxe_standardization_transforms.py, but that doesn't seem to work because this would result in a size mismatch when applying normalization and such. (It seems like dataset_statistics is calculated before std_transform is applied, so you would calculate statistics for a dataset in euler, and then apply it to the transformed quaternion dataset, and thus a size mismatch)

I know in Octo's paper appendix that some experiments regarding training with proprio have been done, so there must be some way to do something like this, right?

Or is using rlds_dataset_mod to manually convert the dataset beforehand the only way to achieve something like this?

Thanks