A Qwen2.5-1.5B-Instruct model fine-tuned with reinforcement learning on BabyAI navigation tasks, reproducing the AgentGym-RL framework (ICLR 2026 Oral, Fudan NLP).
Training Setup
Base model
Qwen/Qwen2.5-1.5B-Instruct
Algorithm
GRPO (Group Relative Policy Optimization)
Environment
BabyAI (GoTo, Pickup, OpenDoor tasks)
Framework
verl + Ray + vLLM 0.6.3 + PyTorch FSDP
Hardware
AWS g4dn.12xlarge (4× NVIDIA T4, 64GB VRAM)
Training steps
31 (checkpoint at step 25)
Batch size
4 episodes per step
The training pipeline uses FSDP to shard the actor model and reference policy across 4 T4 GPUs, with vLLM providing fast rollout generation. The reward signal is the BabyAI task completion score.
What Changed After RL Training
Action grounding (key improvement)
Base model — hallucinates unavailable actions:
Thought: I need to find the red ball first before going to it.
Action: go to red ball 1 ← invalid: ball not in current view
RL-trained model — only selects from the current available actions list:
Thought: The available actions list does not include a direct action to go
to an object without seeing it first.
Action: move forward ← valid: explores until object is visible
The RL model learned that actions must be grounded in the current observation. This is the primary behavioral shift from 31 steps of GRPO training.
Episode trajectory analysis (step 1 vs step 31)
Comparing executor logs at step 1 and step 31 reveals a bimodal pattern that is characteristic of partial training:
Step 1
Step 31
Avg reward
0.64
0.25
Success rate
3/4
1/4
Simple task behavior
Wanders, recovers
1-shot execution (0.991 reward)
Hard task behavior
Wanders aimlessly
"None available" loop (new failure mode)
Interpretation: GRPO successfully reinforced confident, direct action on simple tasks — step 31 can solve a clear pickup task in a single action turn. However, the model has not yet learned to handle ambiguity or unexpected observations gracefully. A new failure mode emerged: when the model is confused, it enters a "None available" loop (repeating that no valid action exists) rather than exploring. This is strictly worse than step 1's naive wandering on hard tasks, and is a classic mid-training artifact — the policy is halfway between random exploration and a converged strategy.
31 steps is intentionally insufficient for full convergence. The goal was to observe the directional behavioral shift, not to reach a production-ready agent.
Multi-Turn Stability
Tested across 4-round simulated BabyAI conversations:
Thought + Action format maintained throughout — no format collapse
Every action selected was from the current available actions list
Thought reasoning occasionally contradictory (expected at 31 steps), but action selection logic was largely sound
Infrastructure Notes
This reproduction involved resolving 16 bugs across the full stack. Key learnings:
T4 compatibility: T4 (sm_75) does not support bfloat16, FlashAttention 2, or Triton bf16 kernels (all require sm_80+). Training uses float16 + SDPA + standard cross-entropy as fallbacks.
verl hybrid engine: load_format=dummy_hf is required for single-node FSDP+vLLM setups. dummy_dtensor causes a deadlock at 43% weight loading.
FSDP checkpoint merging: 4 rank shards → merge along dim=0 → save as safetensors (required due to CVE-2025-32434 blocking .bin loads in newer transformers).
AWS DLAMI: NVMe ephemeral volume clears on instance stop — checkpoints must go on root disk. Account for AdamW optimizer states (~3× model size on disk).
This model expects BabyAI-style prompts with an Available actions: [...] list in each observation turn. It is specifically fine-tuned for action grounding in grid-world navigation tasks.
Limitations
31 training steps is far below convergence. Hard tasks (multi-room, locked doors) are not reliably solved.
The "None available" failure mode on ambiguous observations is a known artifact of partial training.
Task coverage limited to GoTo and Pickup variants encountered during the 31-step run.
Citation
bibtex
1@inproceedings{agentgym2026,
2 title={AgentGym: Evolving Large Language Model-based Agents across Diverse Environments},
3 author={...},
4 booktitle={ICLR 2026},
5 note={Oral presentation}
6}
