VLA-Maze-Solver

---

license: mit tags:

• reinforcement-learning
• robotics
• vision-language-action
• behavior-cloning
• navigation
• pytorch

---

Language-Conditioned Maze Navigation — VLA Agent

A small Vision-Language-Action (VLA) agent that navigates procedurally generated mazes by following free-text instructions (e.g. "go to the red goal", "find the green dot", or even "grass coloured target"), trained via behavior cloning against an A* expert.

Full code, training scripts, and demo videos: GitHub

Model description

The model takes a 64x64 RGB image of the maze and a free-text instruction as input, and outputs a discrete action (up/down/left/right) plus a value estimate. It combines:

• Two frozen CNN encoders — one pretrained to localize agent/goal entities, one pretrained to encode wall/corridor structure
• A frozen CLIP text encoder for the instruction, projected into 4 language tokens
• Cross-attention grounding image patches to the instruction, followed by self-attention for spatial reasoning
• Attention pooling over image patches, feeding into separate action and value heads

See transformer.py in the linked GitHub repo for the full architecture.

Training data

Trained on ~12000+ procedurally generated 5x5 mazes, with an A* planner supplying the optimal action at each sampled state (behavior cloning / imitation learning, not reinforcement learning). Training dataset can be regenerated by running collect_bc_dataset.py from GitHub.

Results

The model generalizes to instruction phrasings never seen during training — including words like "crimson" or "grass coloured" that never appeared in the training set — because the frozen CLIP text encoder already places these semantically close to "red"/"green" in embedding space, and the model's cross-attention learned to key off that direction rather than memorizing literal training sentences.

How to use

import torch
from huggingface_hub import hf_hub_download

# download the checkpoint
ckpt_path = hf_hub_download("Pranav070904/VLA-Maze-Solver", "bc_init_v3_epoch20.pt")

# load into the model (see GitHub repo for the VLA class definition)
from transformer import VLA
agent = VLA()
ckpt = torch.load(ckpt_path, map_location="cpu")
agent.load_state_dict(ckpt["model_state_dict"], strict=False)
agent.eval()

Full demo script (encodes any free-text instruction live with CLIP, generates a fresh maze, saves an episode video) is available in the GitHub repo as demo.py.

Limitations

• Trained and evaluated only on 5x5 procedurally generated mazes with exactly two goal colors (red, green); behavior on substantially larger mazes or additional goal colors is untested.
• Cannot follow negated instructions. Given "don't go to the red goal," the agent still navigates to red. This is a known limitation of CLIP's embedding space — CLIP's contrastive training objective has no explicit mechanism for negation, so "not red" and "red" end up close together rather than opposed. Since the language pathway here is a frozen CLIP encoder, this limitation is inherited rather than something a bigger dataset would fix.

Citation / acknowledgements

Built with CLIP (OpenAI) for language encoding and mazelib for maze generation.