adu2021 Title: Train Agents Efficiently With Synthesized Experience From World Models
Reinforcement learning typically requires millions of real-world interactions, which is expensive and slow. DreamGym solves this by training a reasoning-based experience model that generates synthetic rollouts: step-by-step predictions of state transitions and rewards. These synthetic experiences augment offline real-world data, enabling RL to converge faster with fewer real interactions.
The approach combines three components: experience model (synthesizes dynamics), replay buffer management (balances real and synthetic), and curriculum learning (focuses on useful tasks).
Core Concept
Reasoning-Based Experience Generation for RL:
- Experience Model: Uses step-by-step reasoning to predict state transitions
- Synthetic Rollout Generation: Creates full trajectory sequences without environment interaction
- Hybrid Replay Buffer: Initialized with real data, continuously enriched with synthetic experiences
- Adaptive Curriculum: Actively generates challenging tasks matched to current agent capability
- Efficient RL: Train on diverse synthetic experiences, occasionally validate on real environment
Architecture Overview
- Experience Model: Reasoning-based predictor of (state, action) → (next_state, reward)
- Replay Buffer Manager: Maintains real and synthetic experience ratio
- Curriculum Controller: Determines which synthetic tasks to generate
- Policy Network: Standard RL agent trained on combined data
- Validation Loop: Periodic real-world evaluation to ensure transfer
Implementation Steps
1. Build Reasoning-Based Experience Model
Create a model that generates realistic synthetic experiences using step-by-step reasoning.
class ReasoningBasedExperienceModel(nn.Module):
def __init__(self, state_dim, action_dim, hidden_dim=512):
self.state_encoder = nn.Linear(state_dim, hidden_dim)
self.action_encoder = nn.Linear(action_dim, hidden_dim)
# Transformer for reasoning
self.reasoning_transformer = nn.TransformerEncoder(
nn.TransformerEncoderLayer(d_model=hidden_dim, nhead=8, batch_first=True),
num_layers=6
)
# Output heads
self.next_state_head = nn.Linear(hidden_dim, state_dim)
self.reward_head = nn.Linear(hidden_dim, 1)
self.done_head = nn.Linear(hidden_dim, 1)
def predict_step(self, state, action):
"""Predict (next_state, reward, done) from (state, action)"""
# Encode inputs
state_emb = self.state_encoder(state)
action_emb = self.action_encoder(action)
# Combine with reasoning
combined = state_emb + action_emb
reasoning = self.reasoning_transformer(combined.unsqueeze(1)).squeeze(1)
# Predict outputs
next_state = self.next_state_head(reasoning)
reward = self.reward_head(reasoning)
done = torch.sigmoid(self.done_head(reasoning))
return next_state, reward, done
def generate_trajectory(self, initial_state, policy, max_steps=100):
"""Generate full trajectory using reasoning"""
states = [initial_state]
actions = []
rewards = []
dones = []
state = initial_state
for step in range(max_steps):
# Policy selects action
action = policy.select_action(state)
actions.append(action)
# Experience model predicts outcome
next_state, reward, done = self.predict_step(state, action)
states.append(next_state.detach())
rewards.append(reward.item())
dones.append(done.item() > 0.5)
if done.item() > 0.5:
break
state = next_state
return {
'states': torch.stack(states),
'actions': torch.stack(actions),
'rewards': torch.tensor(rewards),
'dones': torch.tensor(dones)
}
def train(self, real_transitions, learning_rate=1e-4):
"""Train experience model on real transitions"""
optimizer = torch.optim.Adam(self.parameters(), lr=learning_rate)
for state, action, next_state, reward, done in real_transitions:
pred_next, pred_reward, pred_done = self.predict_step(state, action)
# Supervised loss
loss_state = F.mse_loss(pred_next, next_state)
loss_reward = F.mse_loss(pred_reward, reward.unsqueeze(-1))
loss_done = F.binary_cross_entropy(pred_done, done.unsqueeze(-1))
total_loss = loss_state + loss_reward + loss_done
optimizer.zero_grad()
total_loss.backward()
optimizer.step()
2. Implement Hybrid Replay Buffer
Manage mix of real and synthetic experiences.
class HybridReplayBuffer:
def __init__(self, capacity=100000, real_ratio=0.5):
self.real_buffer = []
self.synthetic_buffer = []
self.capacity = capacity
self.real_ratio = real_ratio
def add_real_experience(self, transition):
"""Add actual environment interaction"""
self.real_buffer.append(transition)
if len(self.real_buffer) > self.capacity * self.real_ratio:
self.real_buffer.pop(0)
def add_synthetic_experience(self, trajectory):
"""Add generated experience from model"""
# Convert trajectory to transitions
for i in range(len(trajectory['states']) - 1):
transition = {
'state': trajectory['states'][i],
'action': trajectory['actions'][i],
'reward': trajectory['rewards'][i],
'next_state': trajectory['states'][i + 1],
'done': trajectory['dones'][i]
}
self.synthetic_buffer.append(transition)
# Maintain capacity
max_synthetic = self.capacity * (1 - self.real_ratio)
while len(self.synthetic_buffer) > max_synthetic:
self.synthetic_buffer.pop(0)
def sample_batch(self, batch_size=32):
"""Sample balanced batch from both buffers"""
# Proportion real and synthetic
num_real = int(batch_size * self.real_ratio)
num_synthetic = batch_size - num_real
real_batch = random.sample(self.real_buffer, min(num_real, len(self.real_buffer)))
synthetic_batch = random.sample(self.synthetic_buffer, min(num_synthetic, len(self.synthetic_buffer)))
batch = real_batch + synthetic_batch
random.shuffle(batch)
return {
'states': torch.stack([t['state'] for t in batch]),
'actions': torch.stack([t['action'] for t in batch]),
'rewards': torch.stack([t['reward'] for t in batch]),
'next_states': torch.stack([t['next_state'] for t in batch]),
'dones': torch.stack([t['done'] for t in batch])
}
3. Implement Adaptive Curriculum
Generate synthetic tasks tailored to agent capability.
class AdaptiveCurriculum:
def __init__(self, experience_model, policy, initial_tasks):
self.experience_model = experience_model
self.policy = policy
self.completed_tasks = set()
self.task_difficulty = defaultdict(float)
def estimate_agent_capability(self):
"""Assess current policy strength"""
# Test on standard benchmarks
success_rates = {}
for task in self.completed_tasks:
trajectory = self.experience_model.generate_trajectory(
task['initial_state'], self.policy
)
success = self.evaluate_trajectory(trajectory, task)
success_rates[task] = success
avg_success = np.mean(list(success_rates.values())) if success_rates else 0.5
return avg_success
def generate_curriculum_tasks(self, num_tasks=10):
"""Generate tasks at appropriate difficulty"""
agent_capability = self.estimate_agent_capability()
# Target 70% success rate
target_difficulty = agent_capability + 0.2
new_tasks = []
for _ in range(num_tasks):
# Sample difficulty
difficulty = target_difficulty + np.random.normal(0, 0.1)
difficulty = np.clip(difficulty, 0.0, 1.0)
# Generate task at this difficulty
task = self.generate_task_at_difficulty(difficulty)
new_tasks.append(task)
return new_tasks
def generate_task_at_difficulty(self, difficulty):
"""Procedurally generate task with specified difficulty"""
# Difficulty controls: obstacle density, goal distance, etc.
task = {
'initial_state': self.sample_initial_state(difficulty),
'difficulty': difficulty,
'goal': self.sample_goal(difficulty)
}
return task
4. Integrate Experience Synthesis into RL Training Loop
Combine model learning, experience generation, and policy optimization.
def train_with_experience_synthesis(
experience_model, policy, optimizer,
real_env, initial_buffer,
num_iterations=1000
):
replay_buffer = HybridReplayBuffer()
# Populate with initial real data
for transition in initial_buffer:
replay_buffer.add_real_experience(transition)
curriculum = AdaptiveCurriculum(experience_model, policy, [])
for iteration in range(num_iterations):
# Phase 1: Collect real data (occasionally)
if iteration % 10 == 0:
real_trajectory = collect_real_trajectory(real_env, policy)
for transition in real_trajectory:
replay_buffer.add_real_experience(transition)
# Phase 2: Train experience model on recent real data
real_batch = [t for t in replay_buffer.real_buffer[-100:]]
experience_model.train(real_batch)
# Phase 3: Generate synthetic experiences
curriculum_tasks = curriculum.generate_curriculum_tasks(num_tasks=5)
for task in curriculum_tasks:
synthetic_trajectory = experience_model.generate_trajectory(
task['initial_state'], policy
)
replay_buffer.add_synthetic_experience(synthetic_trajectory)
# Phase 4: Train policy on mixed batch
for _ in range(10):
batch = replay_buffer.sample_batch(batch_size=32)
policy_loss = compute_rl_loss(policy, batch)
optimizer.zero_grad()
policy_loss.backward()
optimizer.step()
# Phase 5: Validation on real environment
if iteration % 100 == 0:
real_eval = evaluate_on_real_env(policy, real_env)
print(f"Iteration {iteration}: Real performance {real_eval:.2%}")
Practical Guidance
When to Use:
- Limited real environment interaction budget
- Expensive real-world rollouts (robotics, simulation)
- Tasks where experience model can be accurate
Hyperparameters:
- real_ratio: 0.5 (balance real and synthetic)
- curriculum_difficulty_step: 0.2 per update
- experience_model_train_freq: Every 10 policy updates
When NOT to Use:
- Environments where model predictions are unreliable
- High-dimensional state spaces (vision) where models struggle
- Tasks requiring perfect memory of environment dynamics
Pitfalls:
- Model drift: Experience model diverges from real environment; validate predictions regularly
- Distribution mismatch: Synthetic data drawn from different distribution than real; use importance weighting
- Compounding errors: Multi-step predictions accumulate errors; keep max trajectory length reasonable
Key Insight: This only works if the experience model is reasonably accurate. Validate model performance on held-out real transitions frequently.