How Rubrics Help AI Agents Learn Across Environments

Peter Bubenik · Apple ML · · Source
Image for RISED: Rubrics for Agentic Multi-Environment Selection and Self-Distillation

Step-by-Step Teaching

Step 1: The Problem Setup — Why Multi-Environment Training Is Hard

Imagine you are training one AI agent to become good at many different tasks simultaneously — like teaching one student to do math, write essays, and code all at once.

The Core Challenge

Environment A (Math)     → Agent learns fast
Environment B (Coding)   → Agent learns slowly  
Environment C (Writing)  → Agent learns medium pace

Problem 1: Uneven learning rates Different environments are mastered at different speeds. This creates two extreme situations within the same training batch:

SituationWhat Happens
All-Failure GroupAgent fails every attempt in Environment B
All-Success GroupAgent succeeds at everything in Environment A

Why is this a problem? Most modern RL training (like GRPO) works by comparing outcomes within a group. If everyone fails or everyone succeeds, there is nothing to compare — no learning signal is generated.


Problem 2: Scalar rewards are informationally poor

A scalar reward is simply a number:

Attempt 1 → reward: 0
Attempt 2 → reward: 0
Attempt 3 → reward: 1

This tells you what happened (success/failure) but not why:

  • Did the agent reason correctly but format wrong?
  • Did it understand the task but make a calculation error?
  • Did it repeat the same mistake three times?

Scalar rewards cannot answer these questions.


Step 2: The Key Insight — Rubrics as Richer Feedback

What Is a Rubric?

A rubric is a structured description of behavior using predefined vocabulary.

Think of it like a teacher's grading checklist:

✅ Positive Rubrics (desired behaviors):
   - "correctly decomposed the problem into sub-steps"
   - "verified the answer before submitting"
   - "used appropriate tool calls"

❌ Negative Rubrics (undesired behaviors):
   - "repeated the same failed approach"
   - "skipped intermediate reasoning"
   - "hallucinated a non-existent function"

Why Rubrics Are More Powerful Than Scalar Rewards

PropertyScalar RewardRubric
Tells you if agent succeeded✅✅
Tells you why it succeeded/failed❌✅
Works when all attempts fail❌✅
Works when all attempts succeed❌✅
Comparable across environments❌✅
Provides token-level guidance❌✅

Key Insight: Rubrics create a shared vocabulary across all environments. Even if Math and Coding tasks look different, both can be tagged with "skipped intermediate reasoning" — enabling cross-environment comparison.


Step 3: The RISED Framework — Three Components

RISED stands for: Rubrics for Agentic Multi-Environment Selection and Self-Distillation

It has three interconnected components:


Component 1: Rubric Tagging (The Foundation)

Rollout (agent's attempt) 
        ↓
   LLM Judge
        ↓
Rubric Profile: {positive tags} + {negative tags}

An LLM judge reads each rollout and assigns rubric tags from a predefined vocabulary shared across all environments.

Example:

Rollout from Math Environment:
"The agent tried 2+2=5, then tried again with 2+2=4"

Rubric Tags:
  Positive: ["self-correction", "persistence"]
  Negative: ["initial arithmetic error"]

This creates a behavioral fingerprint for every rollout.


Component 2: Rubric-Guided Data Selection

Now that every rollout has a rubric profile, RISED uses these profiles to select which data to train on.

The Goal: Select a training batch that:

  1. Reflects the overall behavioral composition of the mixed-environment batch
  2. Limits redundancy — avoids selecting many rollouts with identical rubric profiles

Intuition:

Without rubric selection:
Batch = [Math-fail, Math-fail, Math-fail, Math-fail, Code-success]
         (redundant failures dominate, little diversity)

With rubric selection:
Batch = [Math-fail(type A), Math-fail(type B), Code-success, Write-partial, ...]
         (diverse failure modes + successes = richer learning signal)

This solves the cross-environment relationship problem — selection is now informed by behavioral content, not just which environment the data came from.


Component 3: Self-Distillation Using Rubrics

This component has two directions — one for learning from success, one for avoiding failure:

3a. Positive Rubrics → Self-Distillation Teacher

Successful rollout + its positive rubric tags
              ↓
    Used as "privileged context" 
              ↓
Teacher generates token-level supervision
              ↓
Student (same model) learns more precisely

Why "self-distillation"? The model teaches itself. The positive rubric acts like a hint: "Here is what good behavior looks like — now learn to produce it more reliably."

This provides token-level supervision — not just "this rollout was good" but guidance on which specific tokens/decisions led to success.

3b. Negative Rubrics → Failure Mode Avoidance

Negative rubric tags from current batch
              ↓
Injected into next rollout generation prompt
              ↓
Agent is guided AWAY from recurring failure patterns

Example:

Negative rubrics detected: ["repeated same failed approach", "no tool verification"]

Next generation prompt includes:
"Avoid repeating failed approaches. Always verify tool outputs."

This creates a feedback loop that actively steers future behavior.


Step 4: How the Three Components Work Together

┌─────────────────────────────────────────────────────┐
│                    RISED Pipeline                    │
│                                                     │
│  1. Agent generates rollouts across environments    │
│              ↓                                      │
│  2. LLM Judge assigns rubric tags to each rollout   │
│              ↓                                      │
│  3. Rubric profiles guide DATA SELECTION            │
│     (diverse, representative batch)                 │
│              ↓                                      │
│  4a. Positive rubrics → SELF-DISTILLATION           │
│      (token-level supervision for good rollouts)    │
│                                                     │
│  4b. Negative rubrics → NEXT GENERATION GUIDANCE    │
│      (steer away from failure modes)                │
│              ↓                                      │
│  5. Updated agent generates better rollouts         │
│     → Loop repeats                                  │
└─────────────────────────────────────────────────────┘

Step 5: Why This Generalizes — The Broader Pattern

The paper also mentions two related works that confirm rubric-based approaches work broadly:

RubiCap (Computer Vision)

  • Problem: Dense image captioning is expensive to annotate
  • Solution: Use rubric-guided RL instead of pure supervised distillation
  • Result: Better diversity and generalization than scalar reward RL

Rubric-Based Alignment for QA

  • Problem: Open-domain QA quality has multiple dimensions (accuracy, grounding, clarity)
  • Solution: Generate query-specific rubrics decomposed into quality dimensions
  • Result: Finer-grained supervision than holistic scalar objectives

Pattern: Whenever a task has multiple quality dimensions that scalar rewards collapse into one number, rubrics provide a principled way to recover that structure.


Step 6: Summary — The Core Conceptual Map

PROBLEM                          SOLUTION
─────────────────────────────────────────────────────
All-failure/all-success groups   → Rubrics work without
(no reward contrast)               reward contrast

Cross-environment comparison     → Shared rubric vocabulary
is impossible with scalars         enables comparison

"What went wrong?" is unknown    → Negative rubrics identify
                                   specific failure modes

"How to improve?" is unclear     → Positive rubrics guide
                                   self-distillation

Data selection ignores content   → Rubric profiles enable
                                   content-aware selection

Quick Self-Check Questions

  1. Why do all-failure and all-success rollout groups cause problems for standard RL training?
  2. What makes rubrics "richer" than scalar rewards — name two specific properties?
  3. How does RISED use positive rubrics differently from negative rubrics?
  4. Why is a shared rubric vocabulary across environments important for data selection?
  5. What does "self-distillation" mean in this context, and why is it called that?

More to study