🤖 AI

LaGO: Latent Action Guidance for Online Reinforcement Learning

The paper proposes LaGO, a framework that leverages pretrained Large Language Models as latent action priors to softly guide online reinforcement learning, demonstrating significant improvements in success rates and rewards on both discrete and continuous control benchmarks compared to standard PPO.

Original authors: Kuan-Yen Liu, Ren-Jyun Huang, Ti-Rong Wu

Published 2026-06-24

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Kuan-Yen Liu, Ren-Jyun Huang, Ti-Rong Wu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot how to navigate a complex maze or pick up a specific object. You have two main tools: a brilliant but clumsy expert (a Large Language Model, or LLM) and a stubborn but fast learner (a Reinforcement Learning agent).

For a long time, researchers tried to let the "clumsy expert" drive the robot directly. They asked the expert to say exactly which button to press or which direction to move. The problem? The expert is great at talking and planning, but terrible at giving precise, split-second instructions. If the expert makes a tiny mistake, the robot crashes, and the whole mission fails.

LaGO (Latent Action Guidance) is a new framework that changes the dynamic. Instead of letting the expert drive the car, LaGO lets the expert sit in the passenger seat and whisper directions to the driver.

Here is how it works, broken down into simple steps:

1. The Setup: Two Stages

The paper describes a two-step process to get the robot learning effectively.

Stage 1: The "Study Session" (Offline)
First, the system takes a "clumsy expert" (a pre-trained AI like Llama) and shows it a bunch of videos of humans successfully doing the task (expert demonstrations). It doesn't ask the expert to drive yet. Instead, it teaches the expert to recognize patterns. It's like showing a chess grandmaster a thousand games of chess and asking them to understand the vibe of a good move, rather than forcing them to play a perfect game immediately. The expert learns a "gut feeling" or a latent prior about what actions usually work.
Stage 2: The "Driving Lesson" (Online)
Now, the real learning begins. A separate, fast-learning robot (the RL policy) starts interacting with the real world. It tries things, fails, gets rewards, and learns.
- The Twist: As the robot learns, the "clumsy expert" is still sitting there. But instead of shouting commands, it gently nudges the robot. It says, "Hey, based on what I've seen, you're probably going to want to move this way."
- The robot listens to this nudge but isn't forced to obey. If the robot tries a new move and gets a big reward, it ignores the nudge and keeps going. If the robot is lost, the nudge helps it find a better path.

2. The Analogy: The GPS vs. The Co-Pilot

Think of the old way (using LLMs as direct controllers) as trying to drive a car using a GPS that gives you exact, rigid instructions like "Turn left in 3.42 inches." If the GPS is slightly off, you hit a tree.

LaGO is like having a knowledgeable co-pilot. The co-pilot doesn't grab the steering wheel. Instead, they say, "I think there's a shortcut over there," or "That road looks dangerous."

If the co-pilot is right, you take the shortcut and get there faster.
If the co-pilot is wrong, you ignore them and keep driving your way.
The result? You get the benefit of their experience without the risk of them taking control and crashing the car.

3. The Results: Does it Work?

The researchers tested this on two different "mazes":

CLEVR-Robot: A game where a robot moves balls around on a grid (Discrete steps).
Meta-World: A complex simulation where a robot arm has to open doors, push buttons, and pick up objects (Continuous, smooth movements).

The Outcome:
In both cases, the robots using LaGO learned much better than robots learning alone.

On the ball-moving game, the success rate jumped from 15% to 27%.
On the complex robot arm game, the success rate skyrocketed from 2.7% to 15.2%.

This is a huge deal because the "clumsy expert" wasn't perfect. In fact, if you just let the expert drive the robot directly, it would have failed miserably. But by using the expert as a soft guide, the robot learned to succeed much more often.

4. The Secret Ingredient: A Smarter Expert

The paper also found that the quality of the "clumsy expert" matters.

When they used a smaller, weaker AI as the guide, the robot didn't learn as well.
When they used a bigger, smarter AI (Llama 2 7B vs. TinyLlama), the guidance was much better, and the robot learned even faster.

This suggests that as AI models get smarter in the future, they will automatically become better "co-pilots" for robots, even if they aren't perfect drivers themselves.

Summary

LaGO is a clever way to use the vast knowledge of big AI models to help robots learn, without forcing those models to do the precise, difficult work of controlling the robot. It turns a "know-it-all but clumsy" AI into a helpful guide that speeds up learning and helps the robot succeed more often.

Technical Summary: LaGO – Latent Action Guidance for Online Reinforcement Learning

Problem Statement
While Large Language Models (LLMs) have demonstrated significant potential in planning and sequential decision-making, existing approaches often rely on using them as direct controllers, explicit high-level planners, or world model simulators. These methods require LLMs to generate precise actions, detailed plans, or accurate future trajectories. However, this requirement is often overly strong in practice; LLM capabilities vary significantly, and even small prediction errors can accumulate over time, leading to poor control performance in reinforcement learning (RL) settings that demand precise execution and long-horizon reasoning. Consequently, directly relying on an LLM as the primary controller remains unreliable, particularly in vectorized environments where observations and actions are not naturally represented in language or visual form.

Methodology
To address these challenges, the authors propose LaGO (Latent Action Guidance for Online Reinforcement Learning), a two-stage framework that utilizes a pretrained LLM not as an explicit controller, but as a source of "soft" guidance via latent action priors.

Stage 1: Offline Prior Learning
The framework begins by training a latent policy prior from an offline expert dataset ( $D$ ). A pretrained LLM backbone ( $F_\phi$ ) is used, where non-textual state-action sequences from demonstrations are converted into latent representations via learnable projection modules ( $E_s$ and $E_a$ ). These embeddings, combined with a textual task description ( $x$ ), are fed into the LLM. A newly added action head ( $G_\theta$ ) is then trained via supervised fine-tuning to predict the next action ( $a_{t+1}$ ) given the trajectory prefix ( $\tau_t$ ) and task description. The objective is to minimize the negative log-likelihood of the expert data, creating a latent policy prior $\pi_K(a_{t+1} | x, \tau_t)$ that captures a coarse behavioral prior rather than an optimal controller.
Stage 2: Online RL with Latent Guidance
In the second stage, the trained latent prior $\pi_K$ is frozen and used to guide an online RL policy $\pi_\psi(a|s)$ (e.g., PPO or SAC). The RL agent interacts with the environment to collect rollouts, but the policy optimization is regularized toward the latent prior. The total training loss is defined as:
$L_{online} = L_{PG} + \beta L_{prior}$
where $L_{PG}$ is the standard policy gradient objective and $\beta$ controls the influence of the prior.
- For discrete action spaces, the prior loss is the negative log-likelihood of the prior's predicted action.
- For continuous action spaces, the prior loss is implemented as a cross-entropy term between the prior distribution and the current policy, which serves as an equivalent optimization signal to KL-divergence regularization (up to the entropy of the fixed prior).

This design allows the RL policy to benefit from the structural knowledge encoded in the LLM without requiring the LLM to make exact decisions during interaction.

Key Contributions

Novel Framework: LaGO introduces a method to leverage pretrained LLM knowledge as a latent behavioral prior for online RL, relaxing the strict accuracy requirements of using LLMs as direct controllers.
Two-Stage Architecture: The separation of offline prior learning (via supervised fine-tuning on expert data) and online policy optimization (via regularization) allows the system to utilize LLM knowledge even when the prior is not sufficiently accurate to act as a standalone agent.
Generalization: The framework is designed to work across both discrete and continuous control domains by adapting the prior regularization term to the specific action space.

Experimental Results
The authors evaluated LaGO on two benchmarks: CLEVR-Robot (discrete control) and Meta-World (continuous control). The experiments compared LaGO against "Vanilla PPO" across four task categories: Real, Rephrase, Easy, and Hard.

Performance Gains: LaGO consistently outperformed Vanilla PPO in both reward and success rate.
- On CLEVR-Robot, the average success rate increased from 15.1% to 27.2%.
- On Meta-World, the average success rate increased from 2.7% to 15.2%.
Generalization: Improvements were observed across all task categories, including "Easy" and "Hard" unseen tasks, suggesting the prior provides generalizable guidance rather than mere memorization.
Impact of LLM Quality: An ablation study comparing a Llama-2-7B prior against a TinyLlama-1.1B prior on Meta-World demonstrated that stronger pretrained models yield more effective latent priors. LaGO with Llama-2-7B consistently achieved the highest performance, while the TinyLlama prior sometimes underperformed Vanilla PPO, highlighting the importance of the foundation model's quality.

Significance and Claims
The paper claims that LaGO provides a simple and practical framework for integrating pretrained LLM knowledge into reinforcement learning. The primary significance lies in demonstrating that LLMs can improve online decision-making and planning even when they are not accurate enough to serve as direct controllers. By treating the LLM as a source of soft, latent guidance rather than an explicit planner, the framework mitigates the risks of error accumulation while capitalizing on the rich inductive biases and world knowledge acquired during pretraining.

The authors conclude that the knowledge and generalization abilities encoded in pretrained LLMs can serve as strong priors for planning and RL. They further suggest that as foundation models improve, the effectiveness of this guidance will naturally increase, offering a scalable path for leveraging LLMs in complex decision-making tasks without requiring precise action generation during interaction.