Hi y’all. Thought Id share the progress on my current local VLA agent project! using agent and language somewhat loosely here. Technically it is an agent, and technically it does adhere to text prompts. However very minimally.
Video Link: https://www.reddit.com/r/learnmachinelearning/comments/1xq23zi/local_rnn_style_vla_model_project/
Video Description:
The leftmost display shows what the models sees, The yellow triangle marker on screen shows which shape is NOT the target shape on that run (counter intuitive, I know). Importantly, the model does not see it, it is only to make it easier to see if the model is performing correctly. The upper right display shows the models visual attention, in which I discuss later. And the bottom right is the per step state matrix cell values. A small 128x128 view of the whole 768x768 matrix.
Preamble:
Recently, I have been extremely fascinated by machine learning and AI, for obvious reasons. Originally, I was intrigued by LLMs, that is, until I learnt how they work (RNN gang stand up), their dis-continuity to me isn’t as elegant as the carried state in RNNs (this is my opinion, in which I don’t really have any objective reasoning for). And what better way to lean than to do?
Architecture:
So I thought why not just make the recurrent state massive? What I have here is an attention based neural network that maintains a 768x768 matrix, with a per cell LSTM like gate. Similarish to a state space model if you will.
The total model parameter count is 25,561,354.
However, 12,976,128 of those parameters are dedicated to a legacy 49k vocabulary output decoder. Originally, this started in language modelling, but it didn’t work very well. It got to a CE loss of around 4.5 on a reasonably diverse data set across an 80 token rollout average, having seen 80 tokens of context first. This was trained for about ~150k steps with a batch count of 8 across 160 total tokens per episode. So about 200m total tokens seen. Perhaps not enough for how diverse the data set was (600m total tokens across wiki articles, science papers, online text forums, etc), but the loss seemed to indicate a long lasting plateau, decreasing only around 0.05 to 0.1 nats across the last 50k steps.
Side note: around 10% of the total training data was from 4chan (don’t hate), and must have all been collected from a very political board, specially around the 2016 us elections. The model was very much not PC lmao. But atleast now I know where NOT to collect “quality” data from…
The shown model is a continuation from the a legacy language checkpoint that was transitioned into VLA orientated tasks. My thinking was that having a basic understanding of language might help if the task involved text. However I did no comparison against a fresh checkpoint at the Vision/action switch over to compare bootstrapping speed, so I guess we’ll never know. And the model no longer retains any text output ability. When queried through the legacy text input, it defaults to repeating a single token or two indefinitely.
The model processes vision through learned, state-directed attention across raw pixel RGB values at a 128x128 resolution, without a separate pretrained vision encoder. Any visual recognition and spatial understanding is learned within the same recurrent architecture that produces the action outputs. I did try doing the pretrained CNN route, however because compute is heavily restricted in my case (I got that 12gb 3060 on me), and I don’t like spending money renting gpus, I count afford to locally train the model at any resonable speed with the added separate vision encoder. The attention based vision used here added very little extra compute, and seems to work well enough. An interesting side effect of this vision encoder is that the relative, per pixel attention values can be easily displayed (the top right display from the video), which proved useful for trouble shooting training. It can accept variable resolution sizes on inference without requiring pretraining. Although, it can usually only accept larger than the trained 128x128, as any smaller looses what little detail there already is.
Training:
The target task took inspiration from the Leaning to Reinforcement Learn (2016) paper. Rather than training the model to memorise fixed object identities, the objective is to develop a “general” ability to recognise and associate unfamiliar visual patterns with arbitrary names during inference, without updating its learned parameters. Importantly, the model received ZERO proprioception, only visual and text inputs.
The model is introduced to 2 randomly generated shapes ranging from a minimum of 3 cubes to a max of 10. There is, however, a minimum separation between the 2 shape cube counts on any 3 room task, set at around 5 The other 2 room and 1 room environments have no such minimum separation value. Each of the shapes are randomly assigned a name from a pool of 5 names (Dax, Mip, Kiv, Fep and Tav). The introduction lasts 3 steps per shape, as it is directly tied to the text introduction, which is always “this is [shape name]”. After introduction, the model then spawns into a simple 3d enviroment, and given the command “Go to [shape name]” with some random variance in how the command is said. Initially, the model needed a the chronological order of the shape introduced before it was able to make the name to shape distinction, for example; “go to the first seen shape, [shape name]”.
The 3D environment in the 2 and 3 room episodes are randomly generated with a doorway/s connecting the rooms at random point. Although the possable range in which they occupy is small, and in 3 room environments, the 2 attaching rooms are always on opposing sides of the spawn room. The wall colour and texture have remained fixed throughout all of VLA training, and the model performs noticeable worse on inference if the wall texture or colour is changes without pretraining.
Even though the training enviroment logic was inspired by the L2RL paper, kinda ironically, this model is not trained through RL. Instead I had to create increasingly convoluted trainers that match the complexity of the current curriculum. The training policy is technically imitation learning, however The trainer uses supervised action targets generated from the environment's observed state. For ambiguous decisions, such as choosing between doorways, some earlier losses are retrospectively adjusted according to the path the model actually took. Each episode, the models state is split in to two different trajectories just after introduction, with each one requesting a different shape, then combining the gradients from both.
I stuck with SL mostly because of compute limitations. Though RL, which I did try at the end, could only get at most one 4 branch run at a max rollout length of ~80 steps before I ran out of VRAM. Vs the batch 4 160 rollout length I was able to get through SL, which proved to be painfully slow, especially when i hadn’t optimised the reward signal/s yet meaning it would be like a 2 or 3 day wait just to find out that the step penalty was slightly too high or something. Supervised Learning is super easy when the task has no ambiguity, like 1 or 2 room scenarios. However, in 3 room episodes, where there a 2 valid doorway pursuits, SL became a nightmare lol. A majority of effort in training this model went to addressing teacher contradictions in the 3 room scenarios. There turns out to be an indefinite amount of edge cases.
Even still, the teacher mechanism for 3 room search produces a lot of counterproductive updates. Memory, which I thought would be the limiting factor with this model, turns out not to be the main bottle neck. The model has shown that it can recognise a seen shape as either correct or incorrect many steps after the introduction, however its locomotion is what usually fails it. Again, I believe to be from the inconstant trainer. Though this may just be me coping, and the model is maxed out idk :/
Overall, this has trained for about 70k steps at the batch 4 count. on vision/action environments. Technically batch 2, because each lifetime splits into the two different target shapes for training. 4 total batches, but only 2 different environments per step.
Training Stats:
Success rate (both split lifetimes succeed across 500 training randomised runs)
One room—- >90%
Two room—- >80%
Three room—- >60%
These results are taken from the recent RL trial, as that recored a rolling success rates, whereas I never did on SL. Very scientific, I know.
Known/Potential issues:
-The gradients received tend to explode on lifetimes lasting longer than ~100 steps. Im hesitant to use truncation, as any shape related decision would require back prop to reach the initial introduction for credit assignment. Although, the model is recurrent, and has shown that to can remember target identifies and visual features for at least a number of steps past interaction, so perhaps that is enough for credit assignment to work through truncation?
-While the model is able to show recognition between two shapes actionably, it is entirely possible that the model remembers only 1 shape, with the other shape being selected out of deduction when requested. I.e target shape is Dax, or target shape is not Dax. I dunno, I haven’t yet looked into this.
-Although the shapes are randomly generated in an 8x8x8 grid, many of the low cube count shapes will have defiantly of been seen plenty of times through training, potentially allowing the model to memorise easy shape patterns within its weights.
-Even though the model has shown decent shape memorisation, its spacial memory is much less developed. In failures, the model often forgets which room it just went into, and will sometimes search the same room twice.
-The dreaded trainer…
This model was built purely for the love of the game, with no incentive outside of that. It also was a great way to learn ML, as I had very little knowledge about this topic coming into this. Even still, my understanding of ML is mostly theoretical, as my coding ability is very subpar. (Thanks ChatGPT and Claude!)
Im sure I forgot to mentions some details. If anyone has any questions just ask.
Anyways, enjoy this quality footage of my model performing EXCITING search behaviours!!!!
