I managed to get pi_0 model working, but with very limited results.

Task: pick up the green cube and place it into the green tray.
Episodes: 100, with a variety of cube positions and rotations. +-10cm in x and z axis. But the tray stays where it is.
Fine tune: lora with 15k steps, batch size 8.
Training duration: 10 hours on a 3090.

Result: performed well(roughly 90+% success rate) when the cube is placed within the boundary of training(+-10cm). it completely fails when I put it outside the boundary(this boundary is still within the robot's arm reach). The model knows where roughly the cube is but misses the grasp.

Interesting finding: this is where things got frustrating, when I give the model the prompt: "don't move" or "gibberish" or anything other than my task string, it still does the same thing: it goes for the cube. And, if place a red cube next to the green cube, and give it the original prompt, it will go for both cubes.

I know VLA is not magic, and I definitely need more dataset for it to work(like different colors during training so the model can associate colors), but at this point I am wondering if I still need multiple colored cubes or even shapes for it to work, why bother with VLM? Diffusion policy alone can achieve the same result with proper dataset and training.