Research map / Reinforcement learning / RL from feedback

Controlling Large Language Model with Latent Action

ICML 2025 · Reinforcement learning › RL from feedback · approach: RL for LLM reasoning

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes CoLA, which augments a pre-trained Llama-3.1-8B with a discrete latent action space (inverse dynamics model, action-conditioned language world model, and latent policy) and performs RL over latent actions instead of tokens.

Problem

Token-level RL for LLMs uses a huge action space (vocabulary of 128K+ tokens), making exploration and optimization expensive and prone to knowledge degradation and reward hacking; the paper asks how to learn a compact, well-structured action space for LLM RL.

Approach

Borrowing from reinforcement learning from observations, an inverse dynamics model extracts discrete latent actions from future tokens (codebook of 64 codes, direct code assignment via Gumbel-Softmax); a merge-MLP module injects the action into a pre-trained LLM to form a language world model; a separate policy model outputs latent actions that are trained by BC and RL while the world model is fro

Key results

• CoLA with RL reaches 42.4 on MATH500 vs. 38.2 for the Llama-3.1-8B baseline; with MCTS-Q it reaches 68.2 (baseline MCTS-Q 63.0).
• On Alfworld (unseen) CoLA-RL scores 74.6 (+3.7 over Base-RL 71.6 gain relative to SFT); on Scienceworld seen CoLA-RL scores 28.4 vs. 18.0 for Base-RL.
• CoLA with FTA-P achieves an average GPT-4 win rate of about 64% over the SFT baseline on preference alignment tasks.

Measured effects reported in the paper

Surprising findings

• Baseline RLHF with KL coefficient 0.00 fully failed (reward hacking, degenerate clarifying-question outputs), while CoLA remained robust.
• Full-token-vocabulary MCTS-Q gave the baseline no improvement from the search tricks, whereas CoLA benefited from them.

Limitations

• Only one base model (Llama-3.1-8B) was studied, limiting generality.
• Due to compute limits, the latent action model was trained on only 200G of the 1.1T candidate tokens.
• Agent RL experiments are limited, and the Countdown task showed both models struggled to answer correctly due to base model limitations.

Experimental setup

Datasets: SlimPajama StarCoder Proof-Pile-2 WuDao NuminaMath PRM800k Alpaca-Eval preference data

Benchmarks: MATH500 GSM8K AIME DROP Alfworld Scienceworld Countdown Alpaca-Eval

Metrics: accuracy pass@K semantic diversity (cosine similarity via BGE-M3) GPT-4 win rate score on agent environments format reward

Baselines: Llama-3.1-8B SFT token-level RL baseline MCTS with baseline model MCTS-Q with baseline model

Builds on: VQ-VAE Double DQN GRPO Llama-3.1-8B

Closest papers by meaning

• Reinforcement World Model Learning for LLM-based Agents (NeurIPS 2026 · similarity 0.91)
• Enhancing Decision-Making of Large Language Models via Actor-Critic (ICML 2025 · similarity 0.90)
• Learning Agentic Policy from Action Guidance (NeurIPS 2026 · similarity 0.90)
• Efficient Reinforcement Learning with Large Language Model Priors (ICLR 2025 · similarity 0.90)
• On Predictability of Reinforcement Learning Dynamics for Large Language Models (ICLR 2026 · similarity 0.90)
• Latent Action Reparameterization for Efficient Agent Inference (NeurIPS 2026 · similarity 0.89)