Rwml Reinforcement World Models

---
name: rwml-reinforcement-world-models
title: "Reinforcement World Model Learning for LLM-based Agents"
version: 0.0.2
engine: skillxiv-v0.0.2-claude-opus-4.6
license: MIT
url: "https://arxiv.org/abs/2602.05842"
keywords: [World Models, Self-Supervised Learning, GRPO, Agent Adaptation, Embedding Space]
description: "Train LLM agents to anticipate environment consequences by learning world models through reinforcement learning with embedding-space similarity rewards, avoiding task-specific labels while enabling robust environment adaptation."
---

# Reinforcement World Model Learning for LLM-based Agents

## Problem Context

LLM-based agents excel at language tasks but struggle to anticipate action consequences and adapt to environment dynamics. Standard pretraining emphasizes next-token prediction over semantic understanding of state transitions. This misalignment means agents cannot reliably model how actions transform environments, degrading performance in interactive tasks requiring consequence-aware planning.

## Core Concept

RWML (Reinforcement World Model Learning) uses [GRPO, embedding-space rewards, self-supervised learning] to train agents to predict environment transitions without explicit task rewards or expert annotations. The key insight is measuring prediction quality through semantic similarity in embedding space rather than token-level fidelity, preventing collapse while enabling robust learning.

## Architecture Overview

- **Data collection**: Rollout target model in environments; store interaction traces
- **World model**: Predict next state embedding from current state + action
- **Reward function**: Cosine similarity between predicted and actual next-state embeddings
- **Training**: GRPO optimization on embedding-space rewards; filter challenging samples
- **No task rewards**: Scaling without expert data; world model knowledge transfers to task RL

## Implementation

### Step 1: Collect interaction rollouts

Generate experience by running the base agent in environments. Store state-action-next_state triples with optional natural language annotations.

```python
# Interaction data collection
def collect_rollouts(env, agent, num_episodes=1000, max_steps=50):
    trajectories = []
    for episode in range(num_episodes):
        state = env.reset()
        traj = []
        for step in range(max_steps):
            # Agent generates action
            action = agent.generate_action(state)
            next_state, reward, done = env.step(action)
            traj.append({
                'state': state,
                'action': action,
                'next_state': next_state,
                'reward': reward
            })
            if done:
                break
            state = next_state
        trajectories.append(traj)
    return trajectories
```

### Step 2: Define embedding-space reward function

Compute rewards based on cosine similarity between predicted and actual next-state embeddings. Use frozen encoder from base model.

```python
# Embedding-space reward function
def embedding_reward(predicted_next_state, actual_next_state, encoder):
    """
    Compute cosine similarity in embedding space.
    Prevents token-level collapse while capturing semantic accuracy.
    """
    pred_emb = encoder(predicted_next_state)
    actual_emb = encoder(actual_next_state)

    # Cosine similarity
    similarity = cosine_similarity(
        pred_emb.unsqueeze(0),
        actual_emb.unsqueeze(0)
    ).squeeze()

    # Return similarity as reward (range: [-1, 1], rescale if needed)
    return (similarity + 1.0) / 2.0  # normalize to [0, 1]
```

### Step 3: Filter hard samples for focused learning

Focus training on challenging transitions where the model struggles. Optionally use difficulty-aware weighting.

```python
# Sample filtering: focus on hard negatives
def filter_hard_samples(trajectories, model, encoder, threshold=0.5):
    """
    Keep samples where predicted similarity is below threshold.
    Focus learning on transitions the model finds difficult.
    """
    hard_samples = []
    for traj in trajectories:
        for step in traj:
            pred_next = model.predict_next_state(
                step['state'],
                step['action']
            )
            reward = embedding_reward(
                pred_next,
                step['next_state'],
                encoder
            )

            if reward < threshold:
                hard_samples.append(step)

    return hard_samples
```

### Step 4: Train world model with GRPO

Apply group relative policy optimization with embedding rewards. No task-specific losses needed.

```python
# World model training with GRPO
def rwml_training_step(model, trajectories, encoder, optimizer, lr=1e-5):
    """
    GRPO training on embedding-space rewards for world model.
    """
    # Filter hard samples for focused learning
    hard_samples = filter_hard_samples(trajectories, model, encoder)

    if len(hard_samples) == 0:
        return 0.0

    # Batch processing
    batch_size = 32
    total_loss = 0.0

    for batch_idx in range(0, len(hard_samples), batch_size):
        batch = hard_samples[batch_idx:batch_idx + batch_size]

        # Generate predictions
        predictions = []
        targets = []
        for sample in batch:
            pred = model.predict_next_state(sample['state'], sample['action'])
            predictions.append(pred)
            targets.append(sample['next_state'])

        # Compute rewards using embedding space
        rewards = [
            embedding_reward(pred, target, encoder)
            for pred, target in zip(predictions, targets)
        ]

        # GRPO advantage computation (group relative)
        advantages = compute_grpo_advantages(rewards, group_size=8)

        # Policy gradient step
        loss = -mean(advantages * log_prob(model, batch, predictions))

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_loss += loss.item()

    return total_loss / (len(hard_samples) // batch_size)
```

### Step 5: Combine world model with task RL (optional)

When task-specific rewards available, augment with world model learning for improved performance.

```python
# Combined training: world model + task RL
def combined_training(
    model, env, trajectories, task_reward_fn,
    encoder, optimizer, lambda_world=0.5
):
    """
    Balance world model learning with task-specific RL.
    """
    world_loss = rwml_training_step(model, trajectories, encoder, optimizer)
    task_loss = task_rl_step(model, env, task_reward_fn, optimizer)

    total_loss = lambda_world * world_loss + (1 - lambda_world) * task_loss
    return total_loss
```

## Practical Guidance

**When to use**: Interactive environments (ALFWorld, WebShop, simulators) where predicting environment transitions improves planning. Less beneficial for pure text or static information retrieval tasks.

**Hyperparameters**:
- Embedding similarity threshold for hard sampling (0.4-0.7): lower → more data, higher → focused learning
- GRPO group size (4-8): balance variance and computational cost
- Learning rate (1e-5 to 5e-5): conservative to prevent catastrophic forgetting
- World model weight in combined training (0.3-0.7): higher when environment dynamics are complex

**Common pitfalls**:
- Frozen encoder may not capture task-relevant state properties; consider fine-tuning encoder on domain data
- Embedding similarity doesn't capture all aspects of state (e.g., low-level details); augment with auxiliary losses if needed
- Interaction data distribution matters: diverse trajectories beat large uniform distributions

**Scaling characteristics**: Minimal additional computational overhead; requires only one extra forward pass for next-state encoding. Scales linearly with trajectory count; learning plateaus around 50K-100K well-curated samples.

## Reference

Paper: https://arxiv.org/abs/2602.05842
Related work: GRPO optimization, embedding-based reward functions, world models in RL
Implementation details: Requires environment interaction, off-the-shelf embeddings, standard RL infrastructure