investigate and likely switch to openvla-oft

[mehhl]

there is a new version of OpenVLA, called OpenVLA-OFT link which claims SOTA (as in, better than pi0) performance on variousb enchmarks thanks to enhanced finetuning. The enhancements go far enough to significantly change the model architecture (up to 3 video streams as input, proprio as input, continuous actions rather than tokenization, ...) and the authors are 3 of the original OpenVLA paper authors, one of whom is also a co-founder of Physical Intelligence. It would seem prudent to follow this program:

1. See if OpenVLA-OFT gives some meaningful improvement on our own tasks (nomagic-simple-box, etc.)
2. If the answer to 1. is anything like "yes", switch any/all OpenVLA experiments to OpenVLA-OFT asap.

[mehhl]

For now finetuning our first model (here it is), we'll see in the morning how it behaves. In this one I've basically turned on all the features of OpenVLA-OFT except diffusion, just wanted to see what happens. For now, the loss drops, but with an amplitude reminescent of the very worst parts of reinforcement learning:

Next I should train a sequence (a gradient, if you will) of models from completely vanilla OpenVLA, no FiLM, no nothing but incorporating recent changes to the fine-tuning code, all the way to full OpenVLA-OFT+. Mostly just to see if those changes in finetuning code degrade performance due to bugs. First goal is to reproduce vanilla OpenVLA's best performance with new code.

It also now becomes relevant to implement inference with OpenVLA asap-ish.

[mehhl]

I mean, I don't know. Replay on episode 0:

It is certainly something rather than nothing, but also the model seems to lack courage in some sense, and exhibit frankly bizarre biases (such as on dimension 1). If I had to sum it up in one sentence, the model moves well in some directions but not in others.

I've started a huge rollout of the model over all 230 episodes in the finetuning dataset, it will take about 10 hours. I'd be happy to find out what is the source of model's apparent bias in when to follow demonstrations faithfully vs when to not.

[mehhl]

Another idea is that the dataset might be just a bit too large (230 episodes, of which something like 210 are used for training and 20 for validation). Maybe we should either finetune for longer or on a smaller subset.

[mehhl]

here's another fun visualization:

It plots absolute magnitude of error in the fitted actions $\lvert \hat{a}_{(\cdot)} - a_{(\cdot)}\rvert$, $(\cdot) \in {x, y, z}$ as function of $XYZ$ displacement of the groundtruth action, ie $(a_x, a_y, a_z)$. Here's yet another fun one:

This one plots $\lvert \hat{a}_{(\cdot)} - a_{(\cdot)}\rvert$ against $\lVert (a_x, a_y, a_z) \rVert$.

Here are predicted vs gt acitons for the first 20 episodes:

our problems reside IMO chiefly in the dimension $Z$, or "dimension 2" above. If we were able to make model move confidently downwards, I think that would imply most other problems get solved as well. It is clear that the model has a systemic bias in all three cardinal directions:

which I personally would blame on undertrained action head, then possibly on some unnormalization quirk.

I think the model is undertrained, because there is a certain softness to its errors: they seem like artifacts of some neural process, as evidenced by the fact that the model tends to make similar mistakes in similar places. The artifacts are intricate at the exact scale (5-10 timesteps) I would expect if they came from an undertrained model. Longer clusters of movement (15 timesteps or more) are generally recognized and imitated by the model, but at 5-10 timesteps it didn't have the time to notice most of them yet.

I'll try two things next (in that order):

1. A new training run that records more evaluation info, so that we can see the situation depicted above develop over time. I'll also decrease number of train trajectories to 50 (training on 200 trajs with two A100s for 48h is just silly).

2. A much bigger training run, on the order of 4 A100s for 96h, or 8 A100s for 48h.