SCoRe-Qwen3.5-9B (PLAYPEN, single-player)

Self-correction fine-tune of Qwen/Qwen3.5-9B, trained with a two-stage REINFORCE implementation of SCoRe (Self-Correction via Reinforcement Learning, Kumar et al., 2024) on single-player dialogue games from the PLAYPEN benchmark (Horst et al., EMNLP 2025).

Produced as part of the "Learning in Interaction" seminar project (supervisor: Prof. Sherzod), as well as the Playpen LM Challenge, investigating how SCoRe-style training affects self-correction capability across PLAYPEN's dialogue games.

Model Details

Base model: Qwen/Qwen3.5-9B
Training method: Two-stage SCoRe REINFORCE (as described in arXiv:2409.12917, generalized to multi-turn), 4-bit QLoRA adapter merged into base weights for this release
Training framework: plain HuggingFace transformers + bitsandbytes 4-bit QLoRA + PEFT + 8-bit AdamW, in-process clemcore GameBenchmark (v3.7.2), playpen v3.7.0
Games covered: 5 single-player PLAYPEN dialogue games - adventuregame, textmapworld, textmapworld_graphreasoning, textmapworld_specificroom, wordle. (The remaining 12 PLAYPEN games are multiplayer and require a resolved opponent; deferred to a later phase, not part of this checkpoint.)
License: Apache 2.0 (inherited from base model)

Training Methodology

SCoRe trains a model to improve its own answer across two attempts (y1 → y2) via RL, rewarding genuine self-correction rather than just producing a better first try. This implementation follows the two-stage REINFORCE objective from the reference paper directly, generalized from single-turn Q&A to multi-turn dialogue-game episodes:

Stage I (init): optimize logp(y2) * r(y2); anchor y1 to the reference policy via a KL penalty weighted by beta2. Only y2 receives a policy-gradient signal in this stage.
Stage II (amplify): optimize logp(y1) * r(y1) + logp(y2) * r_shaped(y2), where r_shaped(y2) = r(y2) + alpha * (r(y2) - r(y1)). Both attempts receive a policy-gradient signal and are KL-anchored (weight beta1).

Reference policy: the same weights with LoRA disabled (model.disable_adapter()), computed under torch.no_grad() - no second model is ever resident in GPU memory.

KL estimator: K3 (Schulman) unbiased estimator, exp(d) - d - 1 where d = logp_ref - logp, computed per-turn.

Reward: binary — +1.0 on game success, 0.0 on loss, -1.0 on abort (success_binary mode). Rewards are normalized to [0,1]-scale before the Stage II shaping bonus is applied, so a single alpha=10 is safe and comparable across games with very different native scoring scales.

Length normalization: "constant" mode - policy-gradient and KL terms are divided by max_new_tokens rather than actual generated length, preventing length-gaming and acting as a primary training-stability mechanism.

Environment: games run in-process via clemcore's GameBenchmark API (no clem serve subprocesses, no server ports) — one game instance is played to completion per rollout, keyed by (experiment_name, game_id) to avoid instance-id collisions across experiments that reuse the same numeric ids. This in-process design was adopted specifically to work around server-based OOM and stability issues encountered with an earlier clem serve subprocess architecture.

Generation backend and quantization: the model is loaded via plain transformers + bitsandbytes 4-bit QLoRA (nf4, double-quant, bf16 compute), not Unsloth, and generation uses HF model.generate(), not vLLM. This is a compatibility-driven choice: the Unsloth + vLLM + Qwen3.5-9B combination hit unresolved pip dependency-version conflicts on this cluster stack.

Memory strategy: gradients are accumulated one turn at a time and backpropagated immediately, so peak VRAM is bounded to a single turn's autograd graph rather than growing with the number of turns in a multi-turn game - makes long episodes (e.g. adventuregame) tractable.

Parallel rollout generation: replicas (one per allocated GPU) generate episodes concurrently each round; replica 0 is the sole trainable copy, and its updated LoRA weights are broadcast to the other replicas after each optimizer step before the next round of generation (adapted from Playpen's GRPO-style parallel-generation pattern to SCoRe's two-attempt episodes).

Hyperparameters

Actual values used, from the training SLURM job and confirmed against the training log:

Hyperparameter	Value
alpha (Stage II shaping)	10.0
beta1 (Stage II KL weight, y1 & y2)	0.01
beta2 (Stage I KL weight, y1 only)	0.1
learning rate	1e-5
temperature	0.8
max_new_tokens	256
enable_thinking	false
length normalization	constant (÷ max_new_tokens)
LoRA rank / alpha	16 / 16
LoRA target modules	q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj
LoRA dropout	0.0
quantization	4-bit QLoRA (nf4, double-quant, bf16 compute) via bitsandbytes
optimizer	AdamW 8-bit (bitsandbytes)
stage 1 epochs	2
stage 2 epochs	3
episodes per update (grad accumulation)	8
parallel replicas	4
episode timeout	800s
total optimizer updates (this run)	105
reward mode	success_binary (success=+1, loss=0, abort=−1)
seed	0

Data

Trained on colab-potsdam/playpen-data, restricted to the 5 single-player games listed above. Each epoch plays through the full instance pool per game once (shuffled, not sampled with replacement).

Note: the remaining 12 PLAYPEN games are multiplayer and out of scope for this checkpoint; a multiplayer phase is planned separately.

Compute Budget

Hardware: 3× A30 GPUs
Wall-clock time: ~48 hours

Design Decisions Relevant to Reproducibility

No second resident model for KL: the reference policy is the same weights with LoRA disabled (model.disable_adapter()). Keeps peak VRAM to one model's footprint per replica.
Per-turn backward, not per-episode: gradients are accumulated and backpropagated one dialogue turn at a time (_accumulate_attempt_grads in losses.py) — mathematically equivalent to one combined backward over the whole attempt, but bounds peak memory to a single turn's graph.
use_cache=False is required on the training forward pass: with the KV cache enabled, gradient checkpointing is silently disabled, retaining full-network activations across every turn and causing OOM.
Selective log-softmax avoids materializing a full [T, vocab] log-softmax tensor.
enable_thinking=False is forced for Qwen3.5-9B: it thinks by default, and unsuppressed <think>...</think> content leaks into the game-parser input, aborting episodes.
Instances keyed by (experiment_name, game_id), not game_id alone — adventuregame reuses ids 0-7 across 5 experiments and wordle reuses ids across 2 experiments, so filtering by id alone would silently concatenate unrelated instances into one rollout.
Abort reward = −1.0, not 0.0: a cold-start model aborts most episodes; a zero abort reward deletes the dominant training signal early on.
In-process clemcore.GameBenchmark, not clem serve subprocesses: adopted specifically to resolve server-based OOM and stability problems from an earlier subprocess-per-game architecture.
Plain transformers + bitsandbytes 4-bit QLoRA, not Unsloth: the Unsloth + vLLM + Qwen3.5-9B combination had unresolved pip dependency-version conflicts on this cluster, so load_learner_plain (stock AutoModelForCausalLM, nf4 double-quant, bf16 compute) was used for all replicas instead.
800s episode timeout: long multi-turn games (adventuregame especially) were dropping heavily even at 800s.
Multi-GPU rollout parallelism: 3 A30s, one replica per GPU generating concurrently; replica 0 trains and broadcasts LoRA weights to the others each round.

Known Issues / Limitations

Research artifact for studying self-correction dynamics under RL training on dialogue-game tasks; not intended for production deployment.
Self-correction gains, where present, are specific to the PLAYPEN single-player game distribution and may not transfer to open-ended dialogue or the multiplayer games.

Citation

If you use this model, please cite the SCoRe and PLAYPEN papers:

@misc{kumar2024traininglanguagemodelsselfcorrect,
  title         = {Training Language Models to Self-Correct via Reinforcement Learning},
  author        = {Aviral Kumar and Vincent Zhuang and Rishabh Agarwal and Yi Su and John D Co-Reyes and Avi Singh and Kate Baumli and Shariq Iqbal and Colton Bishop and Rebecca Roelofs and Lei M Zhang and Kay McKinney and Disha Shrivastava and Cosmin Paduraru and George Tucker and Doina Precup and Feryal Behbahani and Aleksandra Faust},
  year          = {2024},
  eprint        = {2409.12917},
  archivePrefix = {arXiv},
  primaryClass  = {cs.LG},
  url           = {https://arxiv.org/abs/2409.12917}
}

@inproceedings{horst2025playpen,
  title     = {Playpen: An Environment for Exploring Learning From Dialogue Game Feedback},
  author    = {Nicola Horst and Davide Mazzaccara and Antonia Schmidt and Michael Sullivan and Filippo Momenté and Luca Franceschetti and Philipp Sadler and Sherzod Hakimov and Alberto Testoni and Raffaella Bernardi and Raquel Fernández and Alexander Koller and Oliver Lemon and David Schlangen and Mario Giulianelli and Alessandro Suglia},
  booktitle = {Proceedings of the 2025 Conference on Empirical Methods in Natural Language Processing},
  pages     = {29854--29891},
  year      = {2025},
  address   = {Suzhou, China},
  publisher = {Association for Computational Linguistics},
  url       = {https://aclanthology.org/2025.emnlp-main.1517/}
}