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
| Metric | Value |
|---|---|
| Validation accuracy (held-out mazes) | 99% |
| Rollout solve rate (unseen mazes, greedy policy) | 96–99% |
| Wrong-goal rate | ~0% |
| Zero-shot instruction generalization (16 unseen phrasings) | 99.2% avg solve rate |
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.