O2 Lab is the research group behind O2 AI. Working with Northwestern University's MLL Lab, we turn frontier research on agentic reasoning, budget-aware agents, and world models into industrial sourcing agents for hardware and electronics supply chains.
Organizing FAGEN @ ICML 2026: Workshop on Failure Modes in Agentic AI
Agents are deployed with token, money, and time budgets. Do they know how much they will spend? BAGEN formalizes budget-awareness as progressive interval estimation: mid-execution, can an agent give a calibrated interval on its remaining budget, and can it tell when a task is no longer finishable?
Fig. 1 — progressive interval estimation. Record a rollout, replay every prefix, and ask the agent: can you still finish, and at what cost?Fig. 2 — solving tasks and predicting their cost are different skills: reward and estimation quality barely correlate across five frontier models.Fig. 10 — the Warehouse environment, contributed by O2 from desensitized enterprise data. The agent runs a manufacturer end-to-end: ordering from OEMs, shipping ocean or air, and filling weekly retailer demand under coupled cost, time, and capacity budgets.
Why it matters for O2: The Warehouse environment models real supply-chain decisions under coupled cost, time, and capacity budgets, the same constrained sourcing decisions O2 automates for electronics teams.
Cite this work
@misc{bagen2026,
title={BAGEN: Are LLM Agents Budget-Aware?},
author={Yuxiang Lin and Zihan Wang and Mengyang Liu and Yuxuan Shan and Longju Bai and Junyao Zhang and Xing Jin and Boshan Chen and Jinyan Su and Xingyao Wang and Jiaxin Pei and Manling Li},
year={2026},
eprint={2606.00198},
archivePrefix={arXiv},
primaryClass={cs.LG},
url={https://arxiv.org/abs/2606.00198},
}
First-authored by Zihan (Zenus) Wang at Northwestern’s MLL Lab.
Identifies “template collapse” in RL-trained multi-turn LLM agents: models produce responses that look diverse but barely depend on the input. The paper decomposes reasoning quality into within-input diversity and cross-input distinguishability via mutual information, explains the collapse through a signal-to-noise-ratio mechanism, and introduces SNR-Aware Filtering to select high-signal prompts.
Mutual-information proxies predict final performance far better than entropy, which is near-useless once collapse starts.
Reasoning length shrinks steadily across eight environments as collapse sets in — a behavioral signature you can watch during training.
SNR-Aware Filtering improves planning, mathematics, web navigation, and code-execution tasks over unfiltered RL.
Fig. 1 — entropy alone can't see the failure. “Template collapse” lives in the top-left quadrant: responses stay diverse, but stop depending on the input.Fig. 4 — SNR-Aware Filtering: rank prompts by within-prompt reward variance and train only on the high-signal ones.
Best Poster Award, Midwest ML Symposium 2025 · 2025-04-24
First-authored by Zihan (Zenus) Wang at Northwestern’s MLL Lab. RAGEN-2 continues this line of work.
The paper behind the RAGEN system: a study of how LLM agents evolve, and fail, when trained with multi-turn reinforcement learning. It introduces StarPO, a trajectory-level RL framework for agents, and uses it to analyze training dynamics across stochastic, multi-step environments.
Introduces StarPO, a general trajectory-level framework for multi-turn agent RL, with RAGEN as the open training system.
Names the Echo Trap, a failure mode where agents converge on repetitive, reward-hacking behavior, and stabilizes training with StarPO-S.
Shows that rollout freshness and diversity are what keep agent self-evolution stable.
Finds that reasoning only emerges when reward design explicitly accounts for it, not from task reward alone.
Co-authored by O2 Tech Lead Zihan (Zenus) Wang, from the same MLL Lab group behind RAGEN.
Trains vision-language agents to reason like world models: before acting, the agent explicitly estimates the current state and predicts what its action will change. A turn-level WorldModeling reward and bi-level credit assignment make this trainable with multi-turn RL across five visual environments.
Frames visual agent training as world-model learning, with explicit state estimation and transition prediction in the reasoning loop.
Introduces a turn-level WorldModeling reward and Bi-Level GAE for turn-aware credit assignment.
A small open model trained with VAGEN outperforms far larger proprietary VLMs on the paper’s multi-turn benchmark suite.
Grounded state reasoning proves more reliable than free-form chain-of-thought for acting agents.
ICLR 2026 · Best Paper, ICCV 2025 SP4V · 2025-06-26
Co-authored by O2 Tech Lead Zihan (Zenus) Wang. Best Paper at the ICCV 2025 SP4V workshop and Best Paper Runner-Up at NeurIPS 2025 LAW.
Can a model build a mental map of a space it has only partially seen? MindCube benchmarks vision-language models on reasoning about unseen parts of a scene from limited views, then shows that having the model first build an internal cognitive map and reason over it substantially beats direct question answering.
Introduces the MindCube benchmark for spatial mental models: reasoning about what is not directly visible from limited camera views.
Finds current VLMs are weak at exactly this, regardless of scale.
A map-then-reason scaffold, where the model constructs a structured cognitive map before answering, delivers large accuracy gains.
Reinforcement learning on top of the scaffold pushes performance further.
Co-authored by O2 Tech Lead Zihan (Zenus) Wang with Stanford collaborators. Oral at the ICCV 2025 LongVid-Foundations workshop.
Finding the few frames that matter in hours of video is a needle-in-a-haystack problem. T* reframes temporal search as spatial search: it lays candidate frames out as image grids and lets a VLM zoom in iteratively, cutting the frames needed for long-video question answering by an order of magnitude.
Introduces LV-Haystack, a large human-annotated benchmark for keyframe search in long videos.
Shows state-of-the-art models locate the right frames far worse than humans do.
T* turns temporal search into iterative spatial reasoning over image grids, with no extra training.
Plugged into existing VLMs, T* lifts long-video QA accuracy while sampling a fraction of the frames.
Co-first-authored by O2 Tech Lead Zihan (Zenus) Wang.
An early, widely cited benchmark for what agents actually do in production: solve tasks over multiple turns, using tools and natural-language feedback. MINT measures how much models improve when they can execute code and receive feedback, and surfaces failure modes that single-turn evaluation hides.
Benchmarks multi-turn task solving with tool use and simulated user feedback, not single-shot answers.
Models gain consistently from tools and feedback, but the gains vary sharply across model families.
Finds that single-turn alignment training (SIFT, RLHF) can hurt multi-turn interaction ability.
Adopted as a standard reference for evaluating interactive LLM agents.
