Rollout Hooks & Async Modes

Slime Integration for Agentic RL Training

← Back to Main README · Overview · Integration Hooks · Generate Functions · Async Modes · Data Source · Multi-Segment · Reward Registry · Slime Wiring

📖 Overview

Dressage hooks into slime via CLI paths — no fork of training code required. The rollout layer bridges slime's RL training loop with Dressage's agentic capabilities through customizable generate functions, async scheduling modes, data sources, and logging hooks. Current slime resolves hook paths as dotted Python import paths, splitting on the last dot.

This design means Dressage rides the slime upgrade train for free: when slime releases new training optimizations, new gradient estimators, or infrastructure improvements, you get them without merge conflicts. Dressage only touches the rollout-side hooks that slime explicitly exposes for customization.

slime training loop
        │
        ├── --custom-generate-function-path     → Agent rollout (blackbox or whitebox)
        ├── --rollout-function-path             → Async scheduling strategy
        ├── --custom-convert-samples-to-...     → Training data conversion
        ├── --custom-rm-path                    → Reward function (sample-oriented)
        ├── --custom-reward-post-process-path   → Advantage broadcast (multi-segment)
        ├── --data-source-path                  → Prompt data loading
        └── --custom-rollout-log-function-path  → Trajectory-level metrics

✨ Integration Hooks

Dressage provides seven hook points that plug into slime's training loop. Each hook handles a specific aspect of the rollout → training pipeline:

🤖 Generate Functions

Two paradigms for agent rollouts, both producing the same list[Sample] output format. The choice between them determines how the agent interacts with the environment, but the post-rollout path (finalize → segments → samples) is identical.

Whitebox Agent

Python agents that call the LLM via proxy and execute tools directly through the paddock. You write the agent logic as a Python class, with full control over the prompting strategy, tool execution, and multi-turn flow.

from dressage.rollout.generate.whitebox_agent import WhiteboxAgent, make_generate

class MyAgent(WhiteboxAgent):
    name = "my_agent"

    async def rollout(self, sample, sampling_params) -> str:
        messages = [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": sample.prompt},
        ]

        response = await self.chat({
            "messages": messages,
            "model": "qwen3.5-4b",
            "tools": self.get_tools(),
        })

        while has_tool_calls(response):
            tool_result = await self.execute_tool(response)
            messages.append(tool_result)
            response = await self.chat({"messages": messages, "model": "qwen3.5-4b"})

        return extract_assistant_content(response)

generate = make_generate(MyAgent)

Agent variants:

• WhiteboxAgent — Pure proxy agent, no sandbox. For lightweight Python tools (API calls, retrieval, computation). Agent calls self.chat(...) to interact with the LLM. Tools are Python functions.
• PaddockWhiteboxAgent — Adds sandbox lifecycle for shell/file execution. Agent calls self.paddock.tool_call(self.session_id, tool_id, args) for isolated command execution. Paddock handles init/terminate automatically.

Blackbox Dispatch

Delegates the entire agent loop to an external HTTP agent running inside a sandbox via BlackboxServer. The generate function manages the lifecycle but doesn't implement any agent logic — that's the backend's job.

1️⃣  paddock.init(traj_id)        → create sandbox lease with BlackboxServer
2️⃣  paddock.register_agent(...)  → start backend process (opencode/openclaw)
3️⃣  paddock.execute_cmd(...)     → run optional setup commands (clone repo, install deps)
4️⃣  paddock.call_agent(...)      → send task prompt, wait for agent completion
5️⃣  proxy.finalize_session(...)  → finalize trajectory → segments
6️⃣  expand_segments_to_samples   → segments → list[Sample]
7️⃣  paddock.terminate(traj_id)   → terminate sandbox lease

Shared Post-Rollout Path

Both whitebox and blackbox modes converge on the same training data pipeline after the agent finishes:

proxy.finalize_session(session_id)
    │  closes all open segments, writes to trajectory store
    ▼
proxy.read_trajectory(session_id)
    │  returns list of finalized segments
    ▼
expand_segments_to_samples(segments, template_sample)
    │  one Sample per segment, shared rollout_id
    │  anchor segment gets reward=None (for reward_fn)
    │  non-anchor segments get reward=0.0
    ▼
list[Sample]
    └── ready for convert_samples_to_train_data

🔄 Async Rollout Modes

Three scheduling strategies for managing rollout batches. The choice determines how much overlap there is between rollout and training, and how latency-tolerant the system is.

Sync Mode (Default)

The simplest mode — slime's default behavior. All samples in a batch must complete before training begins.

Full batch ─── all samples complete ───▶ Train step

• Simplest to understand and debug
• Deterministic batch composition
• Training GPU is idle during the entire rollout phase
• Best for: development, debugging, small-scale experiments

Fully Async Mode

Background worker pipelines keep prompt groups in flight continuously. Completed groups are buffered and assembled into training batches when enough are ready.

Group₁ ─── complete ───▶ ┐
Group₂ ─── complete ───▶ ├─▶ Train step (when enough groups ready)
Group₃ ─── running  ───▶ ┘   (continues in background)
Group₄ ─── running  ───▶     (queued for next train step)
Group₅ ─── queued   ───▶     (waiting for worker)

• Hides blackbox agent latency by overlapping rollout with training
• Flattens list[Sample] from multi-segment generate() automatically
• Better GPU utilization — training can start before all rollouts finish
• More complex scheduling and debugging
• Best for: production blackbox rollouts where agent latency varies widely

Partial Async Mode

Returns early once enough prompt groups fill global_batch_size. Remaining groups continue in background for the next training step. Designed for the common case where rollout_batch_size × n_samples_per_prompt exceeds what training needs.

Groups ─── global_batch_size filled ───▶ Train step
          (e.g., 8 of 16 groups complete)
          remaining groups continue → next step's buffer

• Solves batch-size mismatch efficiently
• No wasted rollout — excess groups carry over to next step
• Natural load balancing across variable-latency rollouts
• Batch composition varies between steps
• Best for: large-scale async rollouts with partial rollout (pause/resume)

# Sync
--rollout-function-path \
  dressage.rollout.sync_rollout.generate_rollout_sync

# Fully async
--rollout-function-path \
  dressage.rollout.fully_async_rollout.generate_rollout_fully_async

# Partial async
--rollout-function-path \
  dressage.rollout.partial_async_rollout.generate_rollout_partial_async

# Override target readiness for partial async
DRESSAGE_PARTIAL_ROLLOUT_TARGET_GROUPS=<int>
DRESSAGE_PARTIAL_ROLLOUT_TARGET_SAMPLES=<int>
DRESSAGE_ROLLOUT_MAX_RETRIES=2
DRESSAGE_ALLOW_EMPTY_TRAIN_BATCH=0

Use slime's sync training entry point (train.py / slime.train) with generate_rollout_sync, slime's async training entry point (train_async.py) with generate_rollout_fully_async, and python3 -m dressage.training.train_async_with_rollout_pause only when proxy pause/resume around weight updates is required.

📂 Data Source

DressageDataSource (dressage/rollout/data_source.py) provides custom prompt data loading for slime. It wraps JSONL prompt files and handles prompt-level metadata needed for prompt-equal aggregation — specifically the instance_id that ties all samples from the same prompt together.

--data-source-path dressage.rollout.data_source.DressageDataSource

Prompt Format

Each line in the JSONL file is a prompt object:

{"prompt": "Fix the authentication bug in auth.py", "label": "tests pass", "metadata": {"instance_id": "inst-001", "reward_fn": "contains_label", "blackbox_execute_cmds": {"before_agent": [{"name": "clone_repo", "cmd": "git clone ...", "required": true, "timeout": 120}]}}}

Sample Datasets

Example prompt datasets are provided in examples/data/:

🧬 Multi-Segment Expansion

Long agentic trajectories often split into multiple segments (due to history compaction, tool schema changes, or TITO fallback). Rather than discarding earlier segments, Dressage trains on all of them via expand_segments_to_samples().

How It Works

Session:  [segment 0] → [segment 1] → [segment 2]
              ↓              ↓              ↓
Samples:   Sample₀        Sample₁        Sample₂  (anchor)
           reward=0.0     reward=0.0     reward=None → reward_fn
           adv broadcast  adv broadcast  adv = A ← terminal reward

Expansion Logic

Core logic in dressage/rollout/multi_segment.py → expand_segments_to_samples():

1. Deep-copy the template sample for each segment
2. Write tokens, masks, logprobs via rollout.artifacts.samples.write_sample_from_segment
3. Tag each sample with parent_traj_id = session ID, segment_index
4. Share the same rollout_id across all segments (grouped in one training step)
5. Anchor segment gets reward=None (filled by reward function); non-anchor segments get reward=0.0
6. Broadcast anchor's advantage to all siblings via reward_post_process

Metrics

Abort Safety

On failure, mark_aborted_no_grad ensures clean cleanup:

• Sets remove_sample=True (excluded from training)
• Stamps parent_traj_id / instance_id for tracking
• Clears session_id for retry availability
• Preserves last_failed_session_id for debugging

📊 Sample Conversion

convert_samples_to_train_data (dressage/rollout/convert_samples.py) transforms Dressage's list[Sample] into slime's train_data format. It is a near-verbatim copy of slime's _convert_samples_to_train_data with one critical addition: prompt-equal rollout_mask_sums for grpo and reinforce_plus_plus_baseline.

Samples → convert_samples_to_train_data → train_data
                                              │
                                              ├── prompt-equal denominators (M_P × N_P / gbs)
                                              ├── rollout_mask_sums per sample
                                              └── fair gradient scaling across segments

🎯 Reward Registry

Pluggable reward functions via dressage/reward/registry.py. Register reward functions with decorators or dynamic module loading, and the registry dispatches them at evaluation time.

from dressage.reward.registry import register_reward

@register_reward("my_custom_reward")
def my_reward(sample, *, args=None):
    """Compute reward for a single sample.

    Args:
        sample: The completed Sample with trajectory data
        args: Optional training args for context

    Returns:
        float: Reward value
    """
    return compute_reward(sample)

Configuration

• Module loading: Reward modules are loaded via DRESSAGE_REWARD_MODULES environment variable (comma-separated module paths)
• Function signature: fn(sample, *, args=None) → float — sample-oriented, receives the complete trajectory
• Slime entry point: dressage.reward.custom_rm.custom_rm — the registry dispatcher
• Selection: sample.metadata["reward_fn"] chooses the registered reward name; missing values fall back to default

# Load custom reward modules
export DRESSAGE_REWARD_MODULES=my_project.rewards,my_project.bonus_rewards

# Wire into slime
--custom-rm-path dressage.reward.custom_rm.custom_rm

📝 Rollout Logging

log_rollout.py computes trajectory-level metrics using parent_traj_id anchor lookup, aggregating across all segments of each trajectory. This gives you a rollout-level view of training progress.

--custom-rollout-log-function-path dressage.rollout.log_rollout.log_rollout_data

🔧 Typical Slime Wiring

A complete example showing how all hooks wire together for a blackbox training run with synchronous rollout and GRPO estimation:

python3 -m slime.train \
  # Generate function: blackbox agent dispatch
  --custom-generate-function-path \
    dressage.rollout.generate.blackbox_dispatch.generate \
  # Sync scheduling: wait for the full rollout batch before training
  --rollout-function-path \
    dressage.rollout.sync_rollout.generate_rollout_sync \
  # Sample conversion: prompt-equal mask sums
  --custom-convert-samples-to-train-data-path \
    dressage.rollout.convert_samples.convert_samples_to_train_data \
  # Reward post-processing: GRPO normalization + segment broadcast
  --custom-reward-post-process-path \
    dressage.training.reward_post_process.reward_post_process \
  # Reward function: pluggable registry
  --custom-rm-path \
    dressage.reward.custom_rm.custom_rm \
  # Data source: JSONL prompt loading
  --data-source-path \
    dressage.rollout.data_source.DressageDataSource \
  # Rollout logging
  --custom-rollout-log-function-path \
    dressage.rollout.log_rollout.log_rollout_data \
  # GRPO advantage estimation
  --advantage-estimator grpo

📊 Upgrade Checklist

When bumping the slime submodule to a new version:

• Diff convert_samples.py against slime rollout.py — check for upstream changes
• Verify build_dp_schedule still uses rollout_id for segment grouping
• Run tests/test_convert_samples_multi_segment.py
• Run tests/test_blackbox_dispatch_multi_segment.py
• Check for new CLI flags or hook points in slime's training loop

📁 Package Structure

dressage/rollout/
├── generate/
│   ├── blackbox_dispatch.py     # Blackbox agent rollout orchestration
│   └── whitebox_agent.py        # WhiteboxAgent + PaddockWhiteboxAgent base classes
├── artifacts/
│   └── samples.py               # Sample writing from trajectory segments
├── fully_async_rollout.py       # Fully async scheduling with worker pipelines
├── partial_async_rollout.py     # Partial async with early return on batch fill
├── sync_rollout.py              # Synchronous rollout (slime default)
├── multi_segment.py             # expand_segments_to_samples()
├── convert_samples.py           # Samples → train_data with prompt-equal masks
├── data_source.py               # DressageDataSource for JSONL prompts
└── log_rollout.py               # Trajectory-level rollout metrics

dressage/reward/
├── custom_rm.py                 # Slime reward entry point (registry dispatcher)
├── registry.py                  # @register_reward decorator + module loading
└── helpers.py                   # Reward computation utilities

🔗 Integration Points

← BlackboxServer · Back to Main README · Next: Training →