Adaptive PRM Simulator

Adaptive Process Reward Model with Generative Difficulty Scaling for LLM Agent Training

An RL training pipeline for LLM agents that co-trains an implicit process reward model with the agent policy, guided by an entropy-based curriculum controller that auto-scales task difficulty.

Architecture

┌─────────────────────────────────────────────────────────┐
│                  AlternatingTrainer                      │
│                                                         │
│  ┌──────────────┐    ┌──────────────┐    ┌───────────┐ │
│  │  Task         │    │  WebShop     │    │ Curriculum │ │
│  │  Generator    │───▶│  Environment │◀───│ Controller │ │
│  │  (difficulty) │    │  (tool-use)  │    │ (entropy)  │ │
│  └──────────────┘    └──────┬───────┘    └─────┬─────┘ │
│                             │                   │       │
│                    trajectories            performance  │
│                             │                   │       │
│  ┌──────────────┐    ┌──────▼───────┐    ┌─────▼─────┐ │
│  │  Implicit     │    │   GRPO       │    │  Reward   │ │
│  │  PRM (DPO)    │───▶│   Trainer    │◀───│  Shaper   │ │
│  │  step rewards │    │  dual-level  │    │ adaptive  │ │
│  └──────────────┘    └──────────────┘    └───────────┘ │
└─────────────────────────────────────────────────────────┘

Key Innovations

1. Implicit Process Rewards via DPO — Instead of expensive MC rollouts, step-level rewards are extracted from a DPO-trained model: r(s_t, a_t) = β · [log π_θ(a_t|s_t) - log π_ref(a_t|s_t)] (inspired by iStar/OPRL)

2. Dual-Granularity GRPO — Extends GRPO with both episode-level (group-relative) and step-level (implicit PRM) advantages, mixed with an adaptive coefficient based on PRM confidence

3. Entropy-Based Curriculum — Three-phase curriculum (discovery → refinement → mastery) with automatic phase transitions driven by performance entropy, inspired by Patronus AI's Generative Simulators

4. Interruption-Robust Credit Assignment — Handles mid-episode perturbations (price changes, stock-outs) with specialized credit assignment strategies: segment-weighted, decay-based, and counterfactual

Project Structure

adaptive-prm-simulator/
├── environments/
│   ├── base_env.py          # Abstract RL environment with tool-use interface
│   ├── webshop_env.py       # WebShop-style shopping agent environment
│   ├── task_generator.py    # Generative task creator with difficulty scaling
│   └── curriculum.py        # Entropy-based curriculum adjuster
├── models/
│   ├── policy.py            # LLM agent policy (Qwen2.5-3B or mock)
│   ├── implicit_prm.py      # Implicit PRM via trajectory DPO
│   └── reward_shaping.py    # Adaptive reward shaping with anti-hacking
├── training/
│   ├── grpo_trainer.py      # GRPO with step+episode level advantages
│   ├── online_prm_update.py # Alternating PRM-policy optimization loop
│   └── interruption_handler.py  # Credit assignment for interrupted episodes
├── evaluation/
│   ├── metrics.py           # Comprehensive metrics suite
│   └── ablations.py         # Systematic ablation experiments
├── configs/
│   └── experiment_configs.yaml
├── main.py                  # Entry point
├── requirements.txt
└── RESEARCH_NOTES.md        # Detailed paper summaries

Quick Start

# Install dependencies
pip install torch numpy pyyaml

# Run smoke test (no GPU required)
python3 main.py --smoke-test

# Run quick training (mock policy)
python3 main.py --config quick_test

# Full training (requires GPU + transformers)
pip install transformers accelerate
python3 main.py --config default --num-steps 5000

Components

Environment Layer

• BaseToolEnv: Abstract interface for tool-use RL environments with trajectory recording and interruption injection
• WebShopEnv: Simulated e-commerce with search/click/filter/buy actions, outcome scoring based on attribute matching, and mid-episode interruptions
• TaskGenerator: Creates shopping tasks of varying difficulty with configurable distractors, attribute complexity, and ambiguity
• CurriculumController: Monitors performance entropy and adjusts difficulty across three phases

Model Layer

• AgentPolicy: Wraps a causal LM as an RL policy with tool-use formatting, log-prob extraction, KL computation, and optional value head
• ImplicitPRM: Learns step-level rewards through DPO on trajectory pairs without explicit step annotations
• AdaptiveRewardShaper: Combines step and episode rewards with horizon adjustment, normalization, and reward hacking detection

Training Layer

• GRPOTrainer: Implements GRPO with dual-granularity advantages (episode-level group-relative + step-level implicit PRM)
• AlternatingTrainer: Orchestrates the full loop — task generation, trajectory collection, PRM scoring, policy update, PRM update, curriculum update
• InterruptionHandler: Segments trajectories at interruption points with three credit assignment strategies

Evaluation Layer

• MetricsTracker: Tracks task completion (overall + by difficulty), sample efficiency, PRM calibration, curriculum progression, and reward hacking rate
• AblationRunner: Six predefined ablation variants for controlled experiments

Ablation Study Design

Variant	Step Rewards	Curriculum	Interruptions	Online PRM
Full System	✓	✓	✓	✓
No Curriculum	✓	✗	✓	✓
No Step Rewards	✗	✓	✓	✓
No Interruptions	✓	✓	✗	✓
Static PRM	✓	✓	✓	✗
Vanilla GRPO	✗	✗	✗	✗

Expected Results

• 10-20% improvement in task completion over vanilla GRPO
• 2-3x sample efficiency gain from process rewards
• Resistance to reward hacking via curriculum adaptation
• Auto-scaling difficulty as agent improves

References

• AgentPRM (2025) — Process Reward Models for LLM Agents
• iStar / OPRL (2025) — Implicit Step Rewards via DPO
• GRPO / DeepSeekMath (2024) — Group Relative Policy Optimization
• Patronus AI Generative Simulators (2025) — Adaptive evaluation environments
• EVO-RAG (2025) — Curriculum-guided RL with discovery/refinement phases