Multi-Domain RL Training

[rafapi]

Enables simultaneous training across multiple domains (math, coding, function calling) with domain-agnostic orchestration

Architecture

Component	Role
multidomain/loader.py	Parses :: syntax, concatenates datasets, injects domain field into each sample
dispatcher.py	Routes problem["domain"] → domain-specific rollout callable via actor.domain_rollouts mapping
domain_sampling.py	Weighted sampling with adaptive rebalancing based on completion ratios to maintain target mix despite varying rollout latencies

Configuration

  actor:
    domain_mix:        # sampling weights (normalised at runtime)
      math: 0.4
      coding: 0.3
      fn_calling: 0.3
    domain_rollouts:   # domain → rollout function mapping
      math: pipelinerl.domains.math.generate_math_rollout
      coding: pipelinerl.domains.coding.generate_coding_rollout
    domain_system_prompts:  # per-domain system prompts
      coding: "You are an expert Python programmer..."

  train_dataset_names:
    - math::open_reasoner_zero_57k
    - coding::coding@train

Adaptive Sampling

DomainWeightedSampler adjusts weights dynamically: adjusted_weight = base_weight × (target_ratio / actual_ratio), clamped to [0.1, 10.0]. This compensates for domains with slower rollouts (e.g. coding sandbox execution) to maintain the configured mix in the output stream.

Metrics

Per-domain stats: domain_mix_actual/{domain}, domain_mix_target/{domain}, domain_mix_count/{domain}.

[ollmer]

Can we have multiple dataloaders at once so we load different datasets for different domains in the same exp?

[rafapi]

yes, it's a proportional concatenation

[ollmer]

so how can we define a multiple dataset_loader functions in a single config?

[rafapi]

in your config you could do (if you wanted to do math, coding and agentic_fn_calling):

defaults:
  - base
  - multi_domain: base        # inits multidomain loader and dispatcher
  - domain_mix: main_mix      # Or inline the mix as showed below

actor:
  domain_mix:
    math: 0.4
    coding: 0.3
    fn_calling: 0.3

train_dataset_names:
  - math::open_reasoner_zero_57k
  - coding::coding@train
  - fn_calling::fn_calling@train

test_dataset_names:
  - math::math_500
  - coding::coding@validation
  - fn_calling::fn_calling@validation

[rafapi]

Domain mix distribution:

Targets: math: 40%, coding: 30%, fn_calling: 30%

Initial distribution:

• math=22.4%
• coding=6.4%
• fn_calling=71.2%

Using adaptive sampling (366k+ completions):

• math=39.6%
• coding=30.3%
• fn_calling=30.1%