Skip to content
SalesforceAIResearchPublic

About

No description, website, or topics provided.

Resources

Code of conduct

Contributing

Security policy

Stars

8 stars

Watchers

0 watching

Forks

πŸ”­ Self-Retrospection Distillation

Prospective Learning: Turning Post-hoc Experiences into Prior Foresight

Haoxiang Zhang*, Qinglin Chen*, Hiroaki Hayashi*, Zhuofeng Li*, Siming Zhang*
Jiaxin Zhang, Jixuan Chen, Fang Wu, Pan Lu, Silvio Savarese, Julian McAuley, Chien-Sheng Wu

Salesforce AI Research Β· UC San Diego Β· Texas A&M University Β· Stanford University

Paper Blog Github


πŸ“° News

TL;DR

Self-Retrospection Distillation (SRD) is an implementation of prospective learning: it aligns the agent's pre-interaction anticipation of what a task and environment will demand with the summary of experience gathered after interacting. As a lightweight auxiliary objective, SRD guides and reinforces RLVR and self-distillation training rather than replacing it. We find that it brings consistent gains over the base GRPO / OPSD / RLSD algorithms across tool-use reasoning, deep research, and simple agentic tasks.

🌟 Overview

SRD overview

Figure 1. SRD pairs a pre-interaction foresight prediction with a hindsight-conditioned one and aligns them on the agent's own foresight rollout β€” turning post-hoc experience into prior foresight.

Contrast with prior experience-use paradigms

SRD vs. prior experience-use paradigms

Figure 2. Retrospective RLVR / self-distillation condition the teacher on hindsight to sharpen the next action; SRD instead uses hindsight to supervise foresight, extracting signal even from reward-uniform groups.

✨ Key features

  • Prospective self-distillation. A lightweight auxiliary objective that aligns the agent's pre-interaction foresight with its post-interaction hindsight, layered on top of RLVR / self-distillation.
  • Model choices from 4B to 35B-A3B. Qwen3.5-4B / 9B / 35B-A3B, one launcher.
  • Diverse eval benchmarks. Math (AIME, AMO-Bench), Code (LiveCodeBench, OJBench), Search (HotpotQA, 2WikiMultiHopQA), and Agent (ALFWorld, WebShop).
  • A sandbox per tool. Each tool β€” code interpreter, search retriever, WebShop, ALFWorld β€” runs in its own containerized sidecar, so they can be deployed separately and on their own resources.
  • Enroot launch, no local setup. A one-click enroot start needs no local environment; unlike Docker it shares the host network namespace with no daemon, adapting cleanly across heterogeneous machines.

πŸ›  Environment & setup

  1. Container image. We use the miles CUDA-12 release image (radixark/miles:v0.1.0-cu12 β€” torch 2.11.0+cu129, with Megatron-LM + SGLang baked in). Point the launcher at it with IMAGE / SQSH (the launcher imports it once into the enroot squashfs).

  2. Install enroot on the host (one-off β€” the image is imported into an enroot squashfs, so no Docker daemon runs inside the training container):

    arch=$(dpkg --print-architecture)
    curl -fSsL -O https://github.com/NVIDIA/enroot/releases/download/v3.5.0/enroot_3.5.0-1_${arch}.deb
    sudo apt install -y ./enroot_3.5.0-1_${arch}.deb
  3. Secrets β€” create examples/SRD/.env (gitignored, never committed):

    WANDB_API_KEY=...          # wandb.ai β†’ User Settings β†’ API keys (wandb.ai/authorize)
    HF_TOKEN=...               # huggingface.co β†’ Settings β†’ Access Tokens
    # optional LLM-as-judge (used only as a search-domain grading fallback):
    OPENAI_API_URL=...         # base URL of your OpenAI-compatible gateway (e.g. https://api.openai.com/v1)
    LLM_GATEWAY_KEY=...        # API key for OPENAI_API_URL (your OpenAI key if that is the endpoint)
    SDPO_REACT_JUDGE_MODEL=gpt-5.6-luna  # judge model name served by that gateway
  4. Scratch root for models / data / checkpoints, on a big disk:

    export SDPO_REACT_LOCAL_ROOT=/path/on/a/big/disk
  5. Data. Training/eval sets are built from public upstreams on first run by the run scripts' [ -f … ] || build_* guards (DAPO-Math-17k, LiveCodeBench, FlashRAG, ALFWorld, WebShop, …). The two pass@k-filtered search-eval datasets (HotpotQA and 2WikiMultiHopQA, 200 prompts) are already prepared in this repo (examples/SRD/data/search_eval/), since they are not reproducible from the builders.

⚑ Quick start

Hardware & versions

GPUs 1 node Γ— 8Γ— H200 (141 GB) recommended. Each run uses all 8 GPUs, so runs are launched sequentially on one node. On 8Γ— H100 or 8Γ— A100-80GB (less VRAM) lower SDPO_ABLATION_MAX_TOKENS_PER_GPU accordingly.
CUDA 12.9 (cu12 image) or 13; NVIDIA driver 570+ for the cu12 build
PyTorch 2.11.0+cu129 (SM90+ / Hopper required for the FlashQLA Qwen GDN attention backend)
Stack Megatron-LM + SGLang + Ray, launched via enroot

Run

πŸ€– One-click for coding agents. Point Claude Code, Codex, or any coding agent at assets/quickstart/SKILL.md β€” a self-contained, step-by-step skill that checks preconditions, writes .env, and launches a run end-to-end (e.g. β€œfollow assets/quickstart/SKILL.md to launch GRPO+SRD on 4B”).

Two reference configs, via the launcher examples/SRD/enroot-run-sdpo-react.sh:

export SDPO_REACT_LOCAL_ROOT=/path/on/a/big/disk

# GRPO + SRD, Qwen3.5-4B, math + code + search
SDPO_REACT_MODEL=qwen3.5-4B SDPO_REACT_RUN_FAMILY=native \
SDPO_ABLATION_ALGO=grpo  SDPO_ABLATION_ARM=e \
  bash examples/SRD/enroot-run-sdpo-react.sh

# OPSD + SRD, Qwen3.5-4B, agentic (ALFWorld + WebShop)
SDPO_REACT_MODEL=qwen3.5-4B SDPO_REACT_RUN_FAMILY=agentic \
SDPO_ABLATION_ALGO=sdpo  SDPO_ABLATION_ARM=e \
  bash examples/SRD/enroot-run-sdpo-react.sh

Arm e = SRD, arm a = no-skill baseline; algo sdpo = OPSD. Swap qwen3.5-4B β†’ 9B / 35B-A3B for other scales. For the detailed algorithm Γ— model matrix β€” every arm, the exact flags, hyperparameters, and env overrides β€” see docs/srd/ablations.md.

Where things land

With DATA_DIR = $SDPO_REACT_LOCAL_ROOT/data/$USER (mounted as /root/data):

  • Rollout / eval dumps β†’ $DATA_DIR/sdpo_dumps/<exp>/
  • Per-trajectory traces (messages, tool calls, grading) β†’ <dump>/agentic_traces/{rollout_id}.jsonl
  • Checkpoints β†’ $DATA_DIR/sdpo_ckpts/ (only when SDPO_REACT_SAVE_CKPT=1)
  • wandb β†’ group sdpo-react-ablation-<config-tag>

πŸ“Š Evaluation

Held-out benchmarks are evaluated during training (--eval-interval 10, avg@8). Pick the suite with SDPO_REACT_EVAL_CONFIG:

  • eval_multitask_min.yaml (default) β€” AIME-2026 + LiveCodeBench-v6-functional + HotpotQA (193 prompts)
  • eval_multitask_full.yaml β€” adds AIME-2024/2025, AMO-Bench, OJBench, 2WikiMultiHopQA

Three standalone harnesses (examples/SRD/ablation/eval-{amo-bench,lcb-functional,ojbench}.sh) report final numbers on a saved checkpoint. For BrowseComp-Plus evaluation of the deep-search agent, please kindly refer to previous work i-DeepSearch/observation-masking.

πŸ“‚ Repository layout

examples/SRD/
β”œβ”€β”€ sdpo.py / reward.py / sdpo_react.py   # method, graders, multi-turn group RM
β”œβ”€β”€ prompt/                               # system / skill / judge prompt text
β”œβ”€β”€ tools/                                # code Β· cli Β· search Β· webshop Β· alfworld sidecars
β”œβ”€β”€ data/                                 # dataset builders + eval configs (+ in-repo search eval)
β”œβ”€β”€ ablation/                             # run scripts (model Γ— domain) + held-out eval harnesses
└── enroot-run-sdpo-react.sh              # one-click launcher
docs/srd/                                 # ablation arms + hyperparameter reference

πŸ”¬ Analysis & Findings

Main results

Main results

Table 1. Main results (avg@8 pass rate, %) on Math, Code, Search, and Agentic. SRD broadly improves GRPO / OPSD / RLSD and transfers beyond training conditions (format shift, unseen ALFWorld, 10Γ—-longer BrowseComp-Plus horizons), mitigates the instability of pure self-distillation, and remains effective as the base policy gets stronger; +SRD rows add SRD on each base algorithm and Improved Perf. rows give the gain.

Sample efficiency

Sample efficiency

Sample efficiency. Across scales, 37–98% of rollout groups are reward-uniform, so group-relative RLVR extracts no gradient from them. SRD keeps learning from exactly these discarded groups β€” e.g. in the 2B all-failure regime (98% uniform) GRPO ends at 0.0% while +SRD reaches 60.6% under the same rollout budget.

SRD displacement

SRD displacement

SRD displacement. LSRD is a dense per-token divergence toward a prospection target derived from the policy itself. Against a host that provides no dense signal (GRPO) it is the only such signal and adds movement; against a host that is already a dense self-distillation (OPSD) it is a second self-referential target competing for the same capacity and behaves like an anchor β€” the run travels less far along the direction it was already going. The plot decomposes the GRPO+SRD update onto GRPO's own direction (blue = kept, red = orthogonal remainder; dashed = unity).

Acknowledgements

Salesforce AI ResearchΒ Β Β  UC San DiegoΒ Β Β  Texas A&M UniversityΒ Β Β  Stanford University

We also thank the following open-source projects:

Citation

@article{zhang2026srd,
  title   = {Self-Retrospection Distillation: Turning Post-hoc Experiences into Prior Foresight},
  author  = {Zhang, Haoxiang and Chen, Qinglin and Hayashi, Hiroaki and Li, Zhuofeng and Zhang, Siming and Zhang, Jiaxin and Chen, Jixuan and Wu, Fang and Lu, Pan and Savarese, Silvio and McAuley, Julian and Wu, Chien-Sheng},
  journal = {arXiv preprint arXiv:2610.08077},
  year    = {2026}
}

Feel free to connect: haz140@ucsd.edu Β· zhuofengli12345@gmail.com Β· hiroakihayashi@salesforce.com Β· wu.jason@salesforce.com

About

No description, website, or topics provided.

Resources

Code of conduct

Contributing

Security policy

Stars

8 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages