In my previous post, I trained an LLM agent on Delta Lake tasks with GRPO and with ECHO. I compared how well each model predicts its tool outputs and how many tasks it solves. That post was already long, so I left out most of what happened on the way to those results. This post covers the problems I ran into, the training metrics that pointed to them, why I read them the way I did, and what I changed.

The setup, briefly:

• Policy: NVIDIA Nemotron 3 Nano 4B, non-thinking. It is a hybrid model that keeps most of its state in Mamba-2 layers and has four attention layers.
• Training: NeMo RL 0.7 on the Megatron backend, with vLLM as the sampler. GRPO with 8 prompts × 8 rollouts per step, for 90 steps.
• Arms: GRPO, ECHO with $\lambda=0.005$ and ECHO with $\lambda=0.02$, one seed each.
• Hardware: each arm ran on one H200 on Modal. A run took about 4 hours and cost about $23.
• Environment: a Databricks-style simulator built on Spark Connect and Delta Lake. The agent has three tools, run_sql, run_python and submit, and a verifier checks the final state of the tables.

The code is at munhouiani/delta_world.

Agentic reinforcement learning usually trains only on the agent’s own tokens: the commands it writes. The tool outputs it reads, such as query results, error messages and table histories, are masked out. The learning signal is one scalar reward per episode.

ECHO (Shrivastava et al., 2026) and related work (PaW) add a second signal. They put a cross-entropy loss on the observation tokens, so the model also learns to predict what the environment will say. They claim for

• faster learning
• better generalization
• less need for expert demonstrations

The intuition is that an agent that knows what a command will do should choose commands better.

I’m writing this post mostly as a reminder to myself. Even though I learned this concept a long time ago, I still find myself coming back to these fundamentals whenever a new training technique comes along.

The Gap

Most tutorials illustrate next-token prediction like the picture on the left: text in, one token out. That’s fine as a mental model, but it doesn’t match what we see in the actual architecture and its outputs.

From PPO to GDPO — Understanding the Evolution of Reinforcement Learning Algorithms

Introduction: Why Reinforcement Learning?

The Two Stages of LLM Training

Training a Large Language Model is fundamentally a two-stage process:

Stage 1: Pre-training (Knowledge Acquisition)

• The model learns from trillions of tokens of text
• It learns to predict the next token: $P(x_t | x_{<t})$
• Result: A “document completor” that understands language structure

Stage 2: Post-training (Behavioral Alignment)

• The model learns to be helpful, harmless, and honest
• It learns to follow instructions and provide useful responses
• Result: A helpful AI assistant

The Problem: A pre-trained model asked “How do I bake a cake?” might respond with “How do I bake a pie?” — continuing a pattern rather than answering the question.

Byte-Pair Encoding

What is Byte-Pair Encoding

Byte-pair encoding (BPE) is a compression algorithm that can be used for text tokenization.

LLaMA 2 Open Foundation and Fine-Tuned Chat Models

📖 Abstract

In this work, we develop and release Llama 2, a collection of pretrained and fine-tuned large language models (LLMs) ranging in scale from 7 billion to 70 billion parameters. Our fine-tuned LLMs, called LLaMA 2-Chat are optimized for dialogue use cases. Our models outperform open-source chat models on most benchmarks we tested, and based on our human evaluations for helpfulness and safety, may be a suitable substitute for closedsource models. We provide a detailed description of our approach to fine-tuning and safety improvements of LLaMA 2-Chat in order to enable the community to build on our work and contribute to the responsible development of LLMs.

LLaMA Open and Efficient Foundation Language Models

📖 Abstract

We introduce LLaMA, a collection of foundation language models ranging from 7B to 65B parameters. We train our models on trillions of tokens, and show that it is possible to train state-of-the-art models using publicly available datasets exclusively, without resorting to proprietary and inaccessible datasets. In particular, LLaMA-13B outperforms GPT-3 (175B) on most benchmarks, and LLaMA65B is competitive with the best models, Chinchilla-70B and PaLM-540B. We release all our models to the research community.

Recall that from Off-Policy Gradient Methods, the objective function is

$$
\begin{aligned}
J^{\theta’}(\theta) &= \hat{\mathbb{E}} _{(s_t, a_t) \sim \pi _{\theta’} (\tau)} \bigg[ \frac{\pi _{\theta}(a_t | s_t)}{\pi _{\theta’}(a_t | s_t)} A^{\theta’}(s_t, a_t) \bigg]
\end{aligned}
$$

We learned that from Importance Sampling, the two distributions shouldn’t differ too much, therefore in PPO, they added a constraint term:

$$
J ^{\theta’} _{PPO} (\theta) = J ^{\theta’} (\theta) - \beta KL(\theta, \theta’)
$$

Where $\beta$ is a weight, this constraint term ensures that $\theta$ and $\theta’$ are similar.

Recall that from Policy Gradient Methods, the calculation of gradient is

$$
\begin{aligned}
\nabla \hat{R} _{\theta} &= \hat{ \mathbb{E}} _{\tau \sim \pi _{\theta}(\tau)} \bigg[ R(\tau) \nabla \log \pi _{\theta} (\tau) \bigg]
\end{aligned}
$$

where

• $\tau$ is a trajectory
• $\pi_{\theta}$ is a stochastic policy
• $\nabla_ {\pi _{\theta}}$ is the gradient operator of the policy $\pi$ with respect to $\theta$
• $\hat{\mathbb{E}} _{\tau \sim \pi _{\theta}(\tau)}[\dots]$ is the empirical average expectation over a finite batch of samples in an algorithm that alternates between sampling and optimization

In this setting, we have to:

• Use $\pi _{\theta}$ to collect data. When the parameter $\theta$ is updated, we must sample the training data again.
• This is called on-policy training

What Are Policy Gradient Methods

Policy gradient methods are a class of reinforcement learning algorithms used for optimizing parametrized policies by gradient descent to maximize the expected return, which is the cumulative reward received over a long-term horizon. These methods are based on updating the policy parameters such that the probabilities of actions leading to higher returns are increased. In comparison, the probabilities of actions leading to lower returns are decreased.

Iteratively optimizing the policy using this approach makes it possible to arrive at an optimal policy that maximizes the expected return. This technique has proven effective in solving complex decision-making problems in various fields, including robotics, game playing, and natural language processing.