Learning On The Job Test Time Curricula

Learning on the Job: Test-Time Curricula for Targeted RL

Core Concept

Traditional training assumes the task distribution is fixed. Learning on the Job (LotJ) flips this assumption: during inference on a target task, the model autonomously assembles a task-specific curriculum from available training data and continues learning via RL. This enables dramatic improvements (1.8-2.1x) on challenging benchmarks by specializing to each problem's characteristics.

Architecture Overview

• Test-Time Curriculum Assembly: Automatically select relevant training examples from large pools without human curation
• Task-Specific Skill Development: Apply RL to continue training on examples most relevant to target task
• Continual Inference-Time Learning: Extend test-time scaling paradigm beyond planning to actual policy updates
• Multi-Domain Generalization: Works across mathematical reasoning (AIME25), competitive coding (CodeElo), diverse task types
• Stateful Learning: Maintain learned skills across examples within same inference session

Implementation Steps

1. Test-Time Curriculum Assembly

Automatically select which training examples to study for each target task.

class TestTimeCurriculumAssembler:
    def __init__(self, training_pool, embedding_model='gpt-4.1'):
        self.training_pool = training_pool  # All available training examples
        self.embedder = embedding_model
        self.selected_curriculum = []

    def assemble_curriculum(self, target_task, curriculum_size=100, budget=50):
        """
        Assemble task-specific curriculum from training pool.

        Args:
            target_task: Target problem to solve
            curriculum_size: Max examples in curriculum
            budget: RL training steps available
        """

        # Step 1: Embed target task
        target_embedding = self.embedder.embed(target_task)

        # Step 2: Retrieve relevant training examples by similarity
        candidate_examples = []
        for example in self.training_pool:
            example_embedding = self.embedder.embed(example['problem'])
            similarity = cosine_similarity(target_embedding, example_embedding)
            candidate_examples.append((similarity, example))

        # Sort by similarity (relevance)
        candidate_examples.sort(reverse=True)

        # Step 3: Select diverse subset
        # Avoid redundancy: pick examples covering different solution patterns
        selected = []
        selected_solutions = set()

        for similarity, example in candidate_examples[:curriculum_size * 2]:
            solution_pattern = self._extract_pattern(example['solution'])

            if solution_pattern not in selected_solutions:
                selected.append(example)
                selected_solutions.add(solution_pattern)

                if len(selected) >= curriculum_size:
                    break

        # Step 4: Order curriculum by difficulty (easy → hard)
        self.selected_curriculum = self._order_by_difficulty(selected, target_task)

        return self.selected_curriculum

    def _extract_pattern(self, solution):
        """Extract solution technique (factorization, recursion, etc.)"""
        # Simplified: could use LLM to extract pattern
        keywords = ['recursion', 'dp', 'binary search', 'greedy', 'factorization']
        patterns = [kw for kw in keywords if kw in solution.lower()]
        return tuple(patterns) if patterns else ('unknown',)

    def _order_by_difficulty(self, examples, target_task):
        """Order examples from easy to hard for curriculum learning."""
        difficulty_scores = []

        for example in examples:
            # Difficulty heuristic: solution length, operation count, etc.
            difficulty = len(example['solution'].split()) / 100  # Normalize
            similarity_to_target = cosine_similarity(
                self.embedder.embed(example['problem']),
                self.embedder.embed(target_task)
            )
            # Easier (but relevant) examples first
            score = difficulty * (1 - similarity_to_target)

            difficulty_scores.append((score, example))

        difficulty_scores.sort()
        return [ex for _, ex in difficulty_scores]

2. Test-Time RL Training Loop

Apply RL to continue learning on selected curriculum examples.

class TestTimeRLTrainer:
    def __init__(self, model, learning_rate=1e-6):
        self.model = model
        self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=learning_rate)

    def train_on_curriculum(self, curriculum_examples, num_steps=50, target_task=None):
        """
        Continue training model on curriculum via RL.
        """

        training_history = []

        for step in range(num_steps):
            # Sample curriculum example
            example = curriculum_examples[step % len(curriculum_examples)]

            # Generate response
            response = self.model.generate(example['problem'])

            # Evaluate correctness
            is_correct = self._evaluate(response, example['solution'])

            # Reward signal
            reward = 1.0 if is_correct else 0.0

            # Policy gradient update
            log_prob = self.model.compute_log_probability(response)

            loss = -log_prob * reward

            # Optional: add entropy regularization
            entropy = self.model.compute_entropy(example['problem'])
            loss = loss - 0.01 * entropy

            # Backward pass
            self.optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
            self.optimizer.step()

            training_history.append({
                'step': step,
                'reward': reward,
                'loss': loss.item(),
                'example_id': example.get('id', step)
            })

        return training_history

    def _evaluate(self, response, reference):
        """Check if response matches reference solution."""
        return extract_answer(response) == extract_answer(reference)

3. Integrated Test-Time Learning Pipeline

Combine curriculum assembly with RL training during inference.

class LearningOnTheJobSystem:
    def __init__(self, model, training_pool, budget_steps=50):
        self.model = model
        self.training_pool = training_pool
        self.budget = budget_steps
        self.trainer = TestTimeRLTrainer(model)
        self.assembler = TestTimeCurriculumAssembler(training_pool)

    def solve_with_learning_on_job(self, target_problem, verbose=False):
        """
        Solve target problem with test-time learning.
        """

        if verbose:
            print(f"Target: {target_problem}")
            print(f"Budget: {self.budget} RL steps\n")

        # Phase 1: Assemble curriculum (task-specific)
        print("Phase 1: Assembling task-specific curriculum...")
        curriculum = self.assembler.assemble_curriculum(
            target_problem,
            curriculum_size=100,
            budget=self.budget
        )
        print(f"Selected {len(curriculum)} relevant training examples\n")

        # Phase 2: Train on curriculum
        print("Phase 2: Learning on the job (test-time RL)...")
        training_history = self.trainer.train_on_curriculum(
            curriculum,
            num_steps=self.budget,
            target_task=target_problem
        )

        # Phase 3: Solve target problem with improved model
        print("Phase 3: Solving target problem...")
        response = self.model.generate(target_problem)

        # Extract answer
        answer = extract_answer(response)

        # Compute learning gains
        avg_reward = sum(h['reward'] for h in training_history) / len(training_history)

        return {
            'answer': answer,
            'reasoning': response,
            'training_history': training_history,
            'avg_curriculum_reward': avg_reward,
            'improvement': f"~{avg_reward:.1%} on curriculum examples"
        }

    def batch_solve(self, target_problems, verbose=False):
        """Solve batch of problems with test-time learning."""
        results = []

        for problem in target_problems:
            result = self.solve_with_learning_on_job(problem, verbose)
            results.append(result)

        return results

4. Benchmark Results

Empirical improvements on AIME25 and CodeElo benchmarks.

# Benchmark improvements
benchmark_results = {
    'AIME25_Math': {
        'qwen_8b_baseline': {
            'pass_at_8': '40%',
            'description': 'Standard sampling without test-time learning'
        },
        'qwen_8b_with_ttc_rl': {
            'pass_at_8': '62%',
            'improvement': '1.8x',
            'curriculum_size': 50,
            'rl_steps': 50
        }
    },
    'CodeElo_Competitive': {
        'qwen_8b_baseline': {
            'pass_at_8': '28%',
            'description': 'Standard generation'
        },
        'qwen_8b_with_ttc_rl': {
            'pass_at_8': '43%',
            'improvement': '2.1x',
            'curriculum_size': 100,
            'rl_steps': 50
        }
    },
    'general_tasks': {
        'improvement_range': '1.3x - 2.5x',
        'factors': [
            'Task difficulty',
            'Training pool quality',
            'Curriculum diversity'
        ]
    }
}

Practical Guidance

Curriculum Assembly: Similarity + diversity is key. High-similarity examples teach task-specific techniques; diverse examples prevent overfitting to specific solution patterns.

RL Budget: 50 steps (~5-10 minutes on modest hardware) yields significant improvements. More steps show diminishing returns (logarithmic scaling).

Training Pool: Larger pools (10K+ examples) enable better curriculum selection. Pool should cover diverse solution approaches for target domain.

Statefulness: Model changes persist across problems within a session. Consider resetting between independent tasks to avoid negative transfer.

When to Use / When NOT to Use

Use When:

• Challenging benchmark problems (math, coding) where specialization helps
• Test-time compute budget available (50-500 RL steps)
• Training pool of relevant examples exists
• Each problem benefits from learning task-specific skills

NOT For:

• Real-time, low-latency inference (RL takes minutes)
• Domains where test-time training causes negative transfer
• Scenarios lacking relevant training examples

Reference

This skill synthesizes findings from "Learning on the Job: Test-Time Curricula for Targeted Reinforcement Learning" (arXiv:2510.04786). Test-time learning extends scaling paradigms beyond planning to active policy improvement during inference.