Professor Wen Ying of Shanghai Jiao Tong University: Building "Generalist" Agents | Agent Insights
Advisor Vitality
One year after our first conversation, we sat down again with Ying Wen, a tenure-track associate professor at Shanghai Jiao Tong University's AI School and John Hopcroft Computer Science Center. Building on his pre-LLM Agent research, Wen has continued to explore Agent research in the large model era, offering perspectives that transcend AI epochs. In this interview, he shared a novel insight: "treat the large model as an Agent, and prompting is essentially having that Agent perform a full online reinforcement learning loop." He and his team proposed an innovative policy optimization method — POAD — applied within multi-agent reinforcement learning frameworks to advance "generalist" Agent research. Today's article focuses on his latest work, POAD. Enjoy.
"Reinforcing Language Agents via Policy Optimization with Action Decomposition"
Paper link: https://arxiv.org/abs/2405.15821
Comparison of action-level Bellman updates (left) versus token-level Bellman updates (right) in POAD
Problem addressed: Language agent actions are typically composed of token sequences (tokens being the minimal symbolic units representing words, phrases, or characters in language models). Directly applying reinforcement learning for policy optimization in language agents generally requires predefining a set of valid actions, while ignoring the fine-grained credit assignment problem within actions at the token level. The team decomposes Agent optimization from the action level down to the token level, providing finer-grained supervision for each token within an action, enabling controllable optimization complexity even in environments with unconstrained language action spaces.
Model architecture: The team introduces Bellman Backup with Action Decomposition (BAD), which integrates intra-action and inter-action token credit assignment processes to provide more efficient, explicit cross-layer credit assignment. They theoretically guarantee that its optimization results remain consistent with traditional pure action-level optimization. By implementing BAD within the PPO algorithm, they introduce Policy Optimization with Action Decomposition (POAD). Benefiting from finer-grained credit assignment and lower optimization complexity, POAD supports policy optimization in environments where the action space is open-ended language, improving learning efficiency and general capabilities of language agents in interactive environments.
Application scenario 1: Embodied intelligence scenarios with constrained action spaces
Application scenario 2: General Agent reinforcement learning framework, mostly used for online learning in open-ended scenarios (data science code generation shown below)
Agent performance after POAD optimization
Agent performance after naive token-level policy optimization (NTPO)
Results: In scenario 1, compared to TWOSOME, POAD achieves equivalent convergence performance, further validating the consistency between BAD's token-level optimization and action-level optimization. Compared to TWOSOME, POAD's training curve is more stable and converges much faster than other baselines — POAD achieves high stability and efficiency through BAD integration. Meanwhile, in scenario 2 (open action space tasks), POAD significantly outperforms naive token-level policy optimization methods, verifying the performance loss caused by NTPO's inconsistency and the gains POAD delivers after ensuring consistency.
*Muning Wen from Shanghai Jiao Tong University is the first author of this paper and also contributed to proofreading this transcript.
Oasis Capital: Could you briefly introduce your research direction in Agents?
Prof. Wen: I previously worked on multi-agent reinforcement learning, which builds on deep reinforcement learning — Agents need to autonomously perceive their environment and take corresponding actions to achieve specific goals. The multi-agent systems we saw in reinforcement learning before were mostly "specialist" Agents: given a task's actions, state space, transition function, and reward function, the Agent learns optimal policies through continuous interaction with the environment. In multi-agent scenarios, beyond environmental interaction, Agents also need to interact with other agents, ultimately learning a strategy capable of cooperating or competing with arbitrary agents.
But "specialists" have issues with sample efficiency and generalization — for instance, an Agent that learns to play Go will struggle to learn StarCraft. This situation persisted until the LLM breakthrough at the end of 2022, when we found that pre-training and fine-tuning on large-scale data could effectively improve model generalization. We then began researching "generalist" Agents, including training decision (action) foundation models from scratch on massive decision trajectory data, or building language/multimodal agents on top of pre-trained language or multimodal foundation models by adding tool use, memory modules, and other components to establish frameworks. These agents continuously interact with environments, other agents, and humans, learning persistently through various feedback signals.
Oasis Capital: Could you explain the original motivation behind developing POAD?
Prof. Wen: Before POAD, we had another paper, "Entropy-Regularized Token-Level Policy Optimization for Large Language Models." At that time, we wanted to establish a unified learning objective for language agents across the pre-training (offline imitation learning) and reinforcement learning online policy optimization phases, enabling continuous improvement in both specialized task performance and general world knowledge (analogous to "reading thousands of books" while "traveling thousands of miles"). But that work wasn't fully mature, so after discussing with doctoral student Muning Wen, we delved deeper into the token-level policy optimization component, which became the POAD paper.
We found that language Agents have enormous action spaces — a single action could range from 3-5 tokens to thousands of tokens, and the action space size grows exponentially with both vocabulary size and the number of tokens composing an action. Previous work mainly considered normalizing actions given a known valid action space. Muning's primary direction is sequence modeling; his previous work MADT (Multi-Agent Decision Transformer) framed multi-agent problems as sequence modeling tasks. So for this problem, we also thought about decomposing language agent actions at the token level and performing sequence modeling to improve policy optimization efficiency. But then we encountered a new problem: naively decomposing actions into token sequences for direct policy optimization effectively optimizes a new Markov decision process, making the optimization problem inconsistent with the original action-level MDP we intended to solve — and this inconsistency causes performance loss in the converged policy. So we began trying to theoretically quantify this gap and explore methods to eliminate it, ultimately arriving at BAD and POAD.
Oasis Capital: What was the focus of your previous work with Muning?
Prof. Wen: Actually, that earlier work had even broader ambitions. We were thinking about how, from a reinforcement learning perspective, large model training is essentially expert imitation learning — all text, characters, and data are human-generated, with next-word prediction as the foundational pre-training objective. If we treat the large model itself as an Agent, and have it engage in online learning during environmental interaction, this becomes interactive reinforcement learning: continuously exploring and sampling policies within the model, evaluating every state and action outcome of the large model Agent based on sampled trajectories or data, then using reinforcement learning for policy improvement to reach optimal policies — a fundamentally different learning objective from imitation learning. Our original paper aimed to integrate these two learning paradigms for large models.
During this integration, we didn't design the unified learning objective particularly well. In the optimization process, we used some token-level optimization methods, which led to the POAD work. We wanted to continuously optimize this token-level optimization process, and found it to be remarkably efficient, so we spun it off into a separate POAD paper.
Oasis Capital: Could you explain POAD's algorithmic structure?
Prof. Wen: POAD's algorithm primarily eliminates the gap between token-level and action-level Bellman equation update processes, enabling consistent policy optimization at the token level with the original action-level MDP. Bellman backup requires computing based on current rewards and discounted future values. When switching directly from action-level Bellman backup to token-level Bellman backup, we go from single-action discount estimation to multiple-token discount estimation, causing inconsistency with original Bellman optimal iteration. We eliminate this inconsistency by distinguishing between inter-action discount factors and intra-action token discount factors, guaranteeing consistency between token-level and action-level policy optimization.
Oasis Capital: Could you explain POAD's evaluation results?
Prof. Wen: Our main evaluations compared against traditional language Agents (GLAM and TWOSOME). We found:
First, POAD provides faster convergence speed and stability, which essentially means higher learning efficiency.
Second, traditional reinforcement learning optimization is constrained by predefined actions — the Agent must know all possible actions in advance. So when optimizing in open spaces, POAD has greater advantages. In scenarios like data science requiring long, open-ended code generation, POAD achieves better results.
A crucial aspect of reinforcement learning research is credit assignment, since rewards are often delayed. Better estimating how actions affect the future is important for policy learning efficiency. By decomposing original actions into finer levels with more precise credit assignment methods, we hope that when given task keywords in the future, Agents can complete them more accurately — preventing individual action tokens from "taking too much blame" or having specific action tokens affected by feedback meant for other tokens.
POAD itself is a general optimization strategy. The advantage of token-level optimization is that in scenarios with sparse rewards, long action sequences, or undefined open-ended scenarios, POAD can more accurately solve policy optimization problems. Going forward, we'll also try embodied scenarios for robot brain-level policy optimization — robot brains essentially perform task decomposition and planning, with each subtask viewable as an action task. We plan to try fine-tuning these tasks precisely to achieve objectives.
Oasis Capital: What are your next steps?
Prof. Wen: We fundamentally do reinforcement learning research. Whether interacting with other Agents or the environment, obtaining feedback from tasks — how to better utilize that feedback for policy learning is a question we've been thinking about.
So, regarding language agents or multimodal agents (note: here referring to agents built on pre-trained language or multimodal foundation models — this is how we think about it):
First, for tasks with quantifiable feedback — where every task has criteria for success or failure — we want to continue exploring how to combine reinforcement learning for more efficient policy optimization. On this foundation, we also hope to extend related research to multi-language-agent cooperation scenarios. For instance, directly applying cooperative multi-agent reinforcement learning algorithms to optimize existing multi-language-agent frameworks (which currently have different manually assigned roles and interaction patterns), or even automatically learning to form different cooperation patterns.
Additionally, we hope to leverage more non-quantifiable feedback from agent interaction processes for continuous learning. In the real world, signals like ratings, clicks, and dwell time — feedback that can't be accurately quantified with numbers — are actually very important. We need to think about how to better collect these signals and incorporate them into the feedback process. Whether to use reward models or other methods for this intermediate transformation is worth considering, ultimately helping large models iterate and learn. But feedback modalities are so diverse (laughs), and quality varies widely — how to define and categorize them is something we've been thinking about. This summer, our group organized a series of AI Agent Workshops, one topic being exploring language agent self-improvement and feedback-driven improvement strategies. We'll share whether we have new findings.
Oasis Capital: What changes do you think large models have brought to Agent and multi-Agent research?
Prof. Wen: Large language models themselves operate at too high an abstraction level. Their generality and generalization appear strong, but they actually struggle to directly solve traditional reinforcement learning tasks requiring fine-grained control. So for professional vertical scenarios, it's difficult to directly apply pre-trained large models — human abstraction of tasks is still needed, combining pre-trained large models with reinforcement learning strategies at higher levels.
However, the pre-training and fine-tuning paradigm of large models has also inspired reinforcement learning training approaches. Drawing on the pre-training思路 of collecting large-scale data and expanding model parameters, followed by online training and learning, can improve generalization in scenarios with relatively similar tasks. But when task differences are substantial, pre-training doesn't learn or generalize well. For example, in robotics scenarios, you can collect extensive expert trajectories for pre-training paths, then do some online learning fine-tuning, considering different levels of action abstraction and using different models and algorithms accordingly. For low-level action scenarios, we use previous reinforcement learning or offline pre-training methods to learn decision/action foundation models; for high-level planning, language foundation model-centric approaches can be used.
Additionally, we can leverage language understanding and scene comprehension capabilities of language models or vision-language models to automatically design better reward functions, write rules for judging what states will be reached, reducing manual effort and making reinforcement learning easier to deploy.
Oasis Capital: How do you view future AI-human collaboration?
Prof. Wen: AI's ultimate goal is to assist humans in completing tasks better. Setting aside AI-dominated modes for now, when AI and humans collaborate as equals on tasks, humans possess a black-box subjective consciousness. In this process, how to enable AI to actively estimate human intentions and purposes, combining AI capabilities to achieve better human-AI coordination, is our current research direction. Of course, this is very difficult for general tasks because tasks are relatively open-ended, and everyone has different subjective ideas — universal human intent inference is hard to achieve. But in specific scenarios, we may be able to analyze different task characteristics to achieve efficient human-machine collaboration, such as machines assisting humans with assembly and similar tasks.
Moreover, AI is increasingly permeating various services and products, but AI isn't omnipotent — in more scenarios, it's AI-assisting-human mode. We need to simultaneously consider AI's strengths and weaknesses, as well as human habits and characteristics, and optimize accordingly to maximize overall utility. For example, humans struggle to quantitatively score data chunks directly, but find comparisons or multiple-choice questions easy. We can use such methods for human preference labeling, then build reward models to train AI alignment.
