Skip to content

Navigation Menu

Sign in
Appearance settings
Sign in
Sign up
Appearance settings
BranchesTags
Open more actions menu

Latest commit

 

History

35 Commits
35 Commits

Folders and files

NameName
Last commit message
Last commit date
verl
verl
 
 
README.md
README.md
 
 

Repository files navigation

Notes to Self: Can LLMs Benefit from Experiential Abstractions?

Pipeline for building a retrieval library of math reasoning abstractions from model traces, injecting them at test time, and optionally fine-tuning with GRPO. All commands run from verl/.

Default stack: MiniLM embeddings (sentence-transformers/all-MiniLM-L6-v2), query rewrite via the base model, and score-gated injection (RETRIEVAL_MODE=score_gate, GATE_MARGIN=0.02).

Environment

git clone https://github.com/moment-timeseries-foundation-model/Test-Time-Training.git

cd Test-Time-Training/verl

conda create -n verl python==3.12
conda activate verl
USE_MEGATRON=0 bash scripts/install_vllm_sglang_mcore.sh
pip install -U "numpy==2.2.0"  # resolve conflicts with Numba
pip install --no-deps -e .

Set OPENROUTER_API_KEY for teacher abstraction extraction (10_extract_full.sh) and abstraction-use judging (judge_abstraction_use.sh). Self-extraction, eval, and GRPO run locally on GPU only.

Datasets

To prepare the training dataset, run:

python -m examples.data_preprocess.math_dataset_ttt --local_save_dir ~/data/math --data_source DigitalLearningGmbH/MATH-lighteval

To prepare MATH-500 for evaluation, run:

python -m examples.data_preprocess.math_dataset_ttt --local_save_dir ~/data/MATH-500 --data_source HuggingFaceH4/MATH-500

The parquet files will be generated under ~/data/.

Pipeline overview

flowchart TD
  trainData[train.parquet] --> collect[00_collect_traces]
  collect --> labeled[traces_labeled.parquet]
  labeled --> teacher[10_extract_full OpenRouter]
  labeled --> selfExt[50_self_extract_full local model]
  teacher --> rawTeacher[raw_abstractions.jsonl]
  selfExt --> rawSelf[raw_self_abstractions.jsonl]
  rawTeacher --> aggregate[20_aggregate MiniLM]
  rawSelf --> aggregateSelf[20_aggregate MiniLM]
  aggregate --> library[library.jsonl]
  aggregateSelf --> librarySelf[library.jsonl]
  trainData --> rewrite[25_rewrite_queries]
  library --> inject[30_inject score_gate]
  librarySelf --> injectSelf[30_inject score_gate]
  rewrite --> inject
  rewrite --> injectSelf
  inject --> eval[40_eval MATH-500]
  injectSelf --> eval
  injectSelf --> grpo[00_grpo naive reward]
  grpo --> merge[merge.sh FSDP to HF]
  merge --> evalTrained[40_eval trained model]
Loading

Common setup

cd verl
conda activate verl

# common exports
export MODEL_FAMILY="meta-llama" #"Qwen"
export MODEL_NAME="Llama-3.2-3B-Instruct" #"Qwen2.5-1.5B-Instruct" #"Qwen3-1.7B-Base"
export MODEL_PATH="$MODEL_FAMILY/$MODEL_NAME"
export TRACE_ROOT="/raid/$USER/traces/$MODEL_NAME"
export EVAL_ROOT="/raid/$USER/eval/abs/$MODEL_NAME"
export CHECKPOINT_ROOT="/raid/$USER/checkpoints/$MODEL_NAME"
export TRAIN_PATH="$HOME/data/math/train.parquet"
export TEST_PATH="$HOME/data/MATH-500/test.parquet"

GPU knobs used by most scripts: CUDA_VISIBLE_DEVICES, NUM_GPUS, GPU_MEM, MICRO_BSZ. These are passed as CLI overrides, so they win over any exps/ config (see "Creating Custom Configs" below).

The hrlib runner scripts each declare their own EXPS stack: 00_grpo.sh uses [grpo,math,hrlib], 40_eval.sh uses [grpo,eval,math,hrlib], and 00_collect_traces.sh / 25_rewrite_queries.sh / 50_self_extract_full.sh use [grpo,eval,math].

Main workflow — teacher extraction (OpenRouter)

Step -1.0: evaluate base model

MODEL_PATH="$MODEL_PATH" \
OUT_DIR="$EVAL_ROOT/base-model/" \
CUDA_VISIBLE_DEVICES="0,1,2,4" \
NUM_GPUS=4 \
GPU_MEM=0.8 \
MICRO_BSZ=2 \
bash examples/hrlib/40_eval.sh

Step -1.1: train and evaluate with GRPO

TRAIN_DATA_PATH="$HOME/data/math/train.parquet" \
MODEL_PATH="$MODEL_PATH" \
OUT_DIR="$CHECKPOINT_ROOT/grpo" \
CUDA_VISIBLE_DEVICES="0,1,2,4" \
NUM_GPUS=4 \
GPU_MEM=0.8 \
MICRO_BSZ=2 \
bash examples/hrlib/00_grpo.sh

LOCAL_DIR="$CHECKPOINT_ROOT/grpo/global_step_58/actor" \
bash examples/test_time_training/merge.sh

MODEL_PATH="$CHECKPOINT_ROOT/grpo/global_step_58/merged_hf_model" \
OUT_DIR="$EVAL_ROOT/grpo/" \
CUDA_VISIBLE_DEVICES="0,1,2,4" \
NUM_GPUS=4 \
GPU_MEM=0.8 \
MICRO_BSZ=1 \
bash examples/hrlib/40_eval.sh

Step 0: collect traces

MODEL_PATH="$MODEL_PATH" \
DATA_PATH="$TRAIN_PATH" \
OUT_DIR="$TRACE_ROOT/round0" \
CUDA_VISIBLE_DEVICES="0,1,2,4" \
NUM_GPUS=4 \
GPU_MEM=0.8 \
MICRO_BSZ=2 \
bash examples/hrlib/00_collect_traces.sh

Step 1: extract abstractions

export OPENROUTER_API_KEY=...

LABELED_PARQUET="$TRACE_ROOT/round0/traces_labeled.parquet" \
OUT_DIR="$TRACE_ROOT/round0" \
MODEL="deepseek/deepseek-v4-flash" \
FALLBACK_MODEL="deepseek/deepseek-v4-flash" \
MAX_CONCURRENCY="40" \
bash examples/hrlib/10_extract_full.sh

Step 2: build MiniLM library

RAW_JSONL="$TRACE_ROOT/round0/raw_abstractions.jsonl" \
EMBEDDER="sentence-transformers/all-MiniLM-L6-v2" \
bash examples/hrlib/20_aggregate.sh

Step 2.1: generate rewritten query parquet

OUT_DIR="$TRACE_ROOT/rewrite_gen" \
IN_PARQUET="$TEST_PATH" \
OUT_PARQUET="$HOME/data/MATH-500/$MODEL_NAME/test_rewritten.parquet" \
DATA_TRAIN="$TRAIN_PATH" \
MODEL_PATH="$MODEL_PATH" \
CUDA_VISIBLE_DEVICES="0,1,2,4" \
NUM_GPUS=4 \
GPU_MEM=0.8 \
MICRO_BSZ=2 \
bash examples/hrlib/25_rewrite_queries.sh

Step 3: inject abstractions

Original query

LIBRARY_DIR="/$TRACE_ROOT/round0" \
IN_PARQUET="$TEST_PATH" \
OUT_PARQUET="$HOME/data/MATH-500/$MODEL_NAME/test_abstraction_orig.parquet" \
QUERY_RECIPE="[{subject}] {user_text}" \
DUMP_SCORES=1 \
RETRIEVAL_MODE=orig \
bash examples/hrlib/30_inject.sh

Rewritten query

LIBRARY_DIR="/$TRACE_ROOT/round0" \
IN_PARQUET="$TEST_PATH" \
OUT_PARQUET="$HOME/data/MATH-500/$MODEL_NAME/test_abstraction_rewrite.parquet" \
QUERY_PARQUET="$HOME/data/MATH-500/$MODEL_NAME/test_rewritten.parquet" \
QUERY_RECIPE="[{subject}] {user_text}" \
DUMP_SCORES=1 \
RETRIEVAL_MODE=rewrite \
bash examples/hrlib/30_inject.sh

Gated retrieval

LIBRARY_DIR="/$TRACE_ROOT/round0" \
IN_PARQUET="$TEST_PATH" \
OUT_PARQUET="$HOME/data/MATH-500/$MODEL_NAME/test_abstraction_gated.parquet" \
QUERY_PARQUET="$HOME/data/MATH-500/$MODEL_NAME/test_rewritten.parquet" \
QUERY_RECIPE="[{subject}] {user_text}" \
DUMP_SCORES=1 \
RETRIEVAL_MODE=score_gate \
GATE_METRIC=top1 \
GATE_MARGIN=0.02 \
GATE_TIE_POLICY=prefer_original \
bash examples/hrlib/30_inject.sh

Evaluation of retrieval stats

python examples/hrlib/score_gate_diagnostics.py \
  --scores "$HOME/data/MATH-500/$MODEL_NAME/meta/test_abstraction_gated_scores.jsonl"

Step 4: evaluate

DATA_VAL="$HOME/data/MATH-500/$MODEL_NAME/test_abstraction_gated.parquet" \
MODEL_PATH="$MODEL_PATH" \
OUT_DIR="$EVAL_ROOT/base-model-inject-gated/" \
CUDA_VISIBLE_DEVICES="0,1,2,4" \
NUM_GPUS=4 \
GPU_MEM=0.8 \
MICRO_BSZ=2 \
bash examples/hrlib/40_eval.sh

Step 5: compare

python3 examples/hrlib/evaluate_results.py lift -- \
--baseline "$EVAL_ROOT/base-model/0.jsonl" \
--treated "$EVAL_ROOT/base-model-inject-gated/0.jsonl"

python3 examples/hrlib/evaluate_results.py lift -- \
--baseline "$EVAL_ROOT/base-model/0.jsonl" \
--treated "$EVAL_ROOT/grpo/0.jsonl"

Step 6: inject abstractions into math train set

Query rewrite

OUT_DIR="$TRACE_ROOT/rewrite_gen_train" \
IN_PARQUET="$TRAIN_PATH" \
OUT_PARQUET="$HOME/data/math/$MODEL_NAME/train_rewritten.parquet" \
DATA_TRAIN="$TRAIN_PATH" \
MODEL_PATH="$MODEL_PATH" \
CUDA_VISIBLE_DEVICES="0,1,2,4" \
NUM_GPUS=4 \
GPU_MEM=0.8 \
MICRO_BSZ=2 \
bash examples/hrlib/25_rewrite_queries.sh

Gated retrieval

LIBRARY_DIR="/$TRACE_ROOT/round0" \
IN_PARQUET="$TRAIN_PATH" \
OUT_PARQUET="$HOME/data/math/$MODEL_NAME/train_abstraction_gated.parquet" \
QUERY_PARQUET="$HOME/data/math/$MODEL_NAME/train_rewritten.parquet" \
QUERY_RECIPE="[{subject}] {user_text}" \
DUMP_SCORES=1 \
RETRIEVAL_MODE=score_gate \
GATE_METRIC=top1 \
GATE_MARGIN=0.02 \
GATE_TIE_POLICY=prefer_original \
bash examples/hrlib/30_inject.sh

Step 7: train and evaluate with GRPO on injected prompts

TRAIN_DATA_PATH="$HOME/data/math/$MODEL_NAME/train_abstraction_gated.parquet" \
MODEL_PATH="$MODEL_PATH" \
OUT_DIR="$CHECKPOINT_ROOT/grpo-injected" \
CUDA_VISIBLE_DEVICES="0,1,2,4" \
NUM_GPUS=4 \
GPU_MEM=0.8 \
MICRO_BSZ=2 \
bash examples/hrlib/00_grpo.sh

LOCAL_DIR="$CHECKPOINT_ROOT/grpo-injected/global_step_58/actor" \
bash examples/test_time_training/merge.sh

MODEL_PATH="$CHECKPOINT_ROOT/grpo-injected/global_step_58/merged_hf_model" \
OUT_DIR="$EVAL_ROOT/grpo-injected/" \
CUDA_VISIBLE_DEVICES="0,1,2,4" \
NUM_GPUS=4 \
GPU_MEM=0.8 \
MICRO_BSZ=1 \
bash examples/hrlib/40_eval.sh

DATA_VAL="$HOME/data/MATH-500/$MODEL_NAME/test_abstraction_gated.parquet" \
MODEL_PATH="$CHECKPOINT_ROOT/grpo-injected/global_step_58/merged_hf_model" \
OUT_DIR="$EVAL_ROOT/grpo-injected-inject-gated/" \
CUDA_VISIBLE_DEVICES="0,1,2,4" \
NUM_GPUS=4 \
GPU_MEM=0.8 \
MICRO_BSZ=2 \
bash examples/hrlib/40_eval.sh

Step 7.5: use Qwen3-1.7B-Base's abstractions

Training

#! need to modify TRAIN_DATA_PATH based on your generation using the above scripts by setting models to Qwen/Qwen3-1.7B-Base
TRAIN_DATA_PATH="$HOME/data/math/Qwen3-1.7B-Base/train_abstraction_re_gated.parquet" \ 
MODEL_PATH="$MODEL_PATH" \
OUT_DIR="$CHECKPOINT_ROOT/grpo-qwen" \
CUDA_VISIBLE_DEVICES="0,1,2,4" \
NUM_GPUS=4 \
GPU_MEM=0.8 \
MICRO_BSZ=2 \
bash examples/hrlib/00_grpo.sh

LOCAL_DIR="$CHECKPOINT_ROOT/grpo-qwen/global_step_58/actor" \
bash examples/test_time_training/merge.sh

MODEL_PATH="$CHECKPOINT_ROOT/grpo-qwen/global_step_58/merged_hf_model" \
OUT_DIR="$EVAL_ROOT/grpo-injected-qwen/" \
CUDA_VISIBLE_DEVICES="0,1,2,4" \
NUM_GPUS=4 \
GPU_MEM=0.8 \
MICRO_BSZ=1 \
bash examples/hrlib/40_eval.sh

DATA_VAL="$HOME/data/MATH-500/$MODEL_NAME/test_abstraction_gated.parquet" \
MODEL_PATH="$CHECKPOINT_ROOT/grpo-qwen/global_step_58/merged_hf_model" \
OUT_DIR="$EVAL_ROOT/grpo-injected-qwen-inject-gated/" \
CUDA_VISIBLE_DEVICES="0,1,2,4" \
NUM_GPUS=4 \
GPU_MEM=0.8 \
MICRO_BSZ=2 \
bash examples/hrlib/40_eval.sh

Step 8: compare

python3 examples/hrlib/evaluate_results.py lift -- \
--baseline "$EVAL_ROOT/base-model/0.jsonl" \
--treated "$EVAL_ROOT/base-model-inject-gated/0.jsonl"

python3 examples/hrlib/evaluate_results.py lift -- \
--baseline "$EVAL_ROOT/grpo/0.jsonl" \
--treated "$EVAL_ROOT/grpo-injected/0.jsonl"

python3 examples/hrlib/evaluate_results.py lift -- \
--baseline "$EVAL_ROOT/base-model-inject-gated/0.jsonl" \
--treated "$EVAL_ROOT/grpo-injected-inject-gated/0.jsonl"

On Qwen abstractions

python3 examples/hrlib/evaluate_results.py lift -- \
--baseline "$EVAL_ROOT/grpo/0.jsonl" \
--treated "$EVAL_ROOT/grpo-qwen/0.jsonl"

python3 examples/hrlib/evaluate_results.py lift -- \
--baseline "$EVAL_ROOT/grpo-qwen/0.jsonl" \
--treated "$EVAL_ROOT/grpo-qwen-inject-gated/0.jsonl"

Cross eval (on prompts injected with Qwen)

python3 examples/hrlib/evaluate_results.py lift -- \
--baseline "$EVAL_ROOT/base-model/0.jsonl" \
--treated "$EVAL_ROOT/base-model-cross/0.jsonl"

python3 examples/hrlib/evaluate_results.py lift -- \
--baseline "$EVAL_ROOT/grpo-injected-inject-gated/0.jsonl" \
--treated "$EVAL_ROOT/grpo-injected-cross/0.jsonl"

python3 examples/hrlib/evaluate_results.py lift -- \
--baseline "$EVAL_ROOT/grpo-injected-inject-gated/0.jsonl" \
--treated "$EVAL_ROOT/grpo-qwen-cross/0.jsonl"

Judge abstraction use

MODEL="deepseek/deepseek-v4-flash" \
FALLBACK_MODEL="deepseek/deepseek-v4-flash" \
VAL_JSONL="$EVAL_ROOT/base-model-inject-gated/0.jsonl" \
OUT_DIR="$EVAL_ROOT/base-model-inject-gated/judge-abs-use" \
bash examples/hrlib/judge_abstraction_use.sh

python examples/hrlib/evaluate_results.py judge-summary -- \
  "$EVAL_ROOT/base-model-inject-gated/judge-abs-use/judge_results.jsonl"

python examples/hrlib/evaluate_results.py judge-summary -- \
  "$EVAL_ROOT/grpo-injected-inject-gated/judge-abs-use/judge_results.jsonl"

Edit global_step_58 in merge commands to match your checkpoint step.

Self extraction — base model extracts its own abstractions

What if we use the base model itself to extract abstractions?

Collect traces

MODEL_PATH="$MODEL_PATH" \
DATA_PATH="$TRAIN_PATH" \
OUT_DIR="$TRACE_ROOT/round0" \
CUDA_VISIBLE_DEVICES="0,1,2,4" \
NUM_GPUS=4 \
GPU_MEM=0.8 \
MICRO_BSZ=2 \
bash examples/hrlib/00_collect_traces.sh

Self-extract abstractions

LABELED_PARQUET="$TRACE_ROOT/round0/traces_labeled.parquet" \
DATA_TRAIN="$HOME/data/math/train.parquet" \
MODEL_PATH="$MODEL_PATH" \
CUDA_VISIBLE_DEVICES="0,1,2,4" \
NUM_GPUS=4 \
GPU_MEM=0.8 \
MICRO_BSZ=2 \
bash examples/hrlib/50_self_extract_full.sh

Build MiniLM library

RAW_JSONL="$TRACE_ROOT/round0/raw_self_abstractions.jsonl" \
OUT_DIR="$TRACE_ROOT/round0/self" \
EMBEDDER="sentence-transformers/all-MiniLM-L6-v2" \
bash examples/hrlib/20_aggregate.sh

Rewrite train and test

OUT_DIR="$TRACE_ROOT/rewrite_gen_train" \
IN_PARQUET="$TRAIN_PATH" \
OUT_PARQUET="$HOME/data/math/$MODEL_NAME/train_rewritten.parquet" \
DATA_TRAIN="$TRAIN_PATH" \
MODEL_PATH="$MODEL_PATH" \
CUDA_VISIBLE_DEVICES="0,1,2,4" \
NUM_GPUS=4 \
GPU_MEM=0.8 \
MICRO_BSZ=2 \
bash examples/hrlib/25_rewrite_queries.sh

OUT_DIR="$TRACE_ROOT/rewrite_gen" \
IN_PARQUET="$TEST_PATH" \
OUT_PARQUET="$HOME/data/MATH-500/$MODEL_NAME/test_rewritten.parquet" \
DATA_TRAIN="$TRAIN_PATH" \
MODEL_PATH="$MODEL_PATH" \
CUDA_VISIBLE_DEVICES="0,1,2,4" \
NUM_GPUS=4 \
GPU_MEM=0.8 \
MICRO_BSZ=2 \
bash examples/hrlib/25_rewrite_queries.sh

Gated inject on test

LIBRARY_DIR="/$TRACE_ROOT/round0/self" \
IN_PARQUET="$TEST_PATH" \
OUT_PARQUET="$HOME/data/MATH-500/$MODEL_NAME/test_self_abstraction_gated.parquet" \
QUERY_PARQUET="$HOME/data/MATH-500/$MODEL_NAME/test_rewritten.parquet" \
QUERY_RECIPE="[{subject}] {user_text}" \
DUMP_SCORES=1 \
RETRIEVAL_MODE=score_gate \
GATE_METRIC=top1 \
GATE_MARGIN=0.02 \
GATE_TIE_POLICY=prefer_original \
bash examples/hrlib/30_inject.sh

Gated inject on train

LIBRARY_DIR="/$TRACE_ROOT/round0/self" \
IN_PARQUET="$TRAIN_PATH" \
OUT_PARQUET="$HOME/data/math/$MODEL_NAME/train_self_abstraction_gated.parquet" \
QUERY_PARQUET="$HOME/data/math/$MODEL_NAME/train_rewritten.parquet" \
QUERY_RECIPE="[{subject}] {user_text}" \
DUMP_SCORES=1 \
RETRIEVAL_MODE=score_gate \
GATE_METRIC=top1 \
GATE_MARGIN=0.02 \
GATE_TIE_POLICY=prefer_original \
bash examples/hrlib/30_inject.sh

Eval base

DATA_VAL="$HOME/data/MATH-500/$MODEL_NAME/test_self_abstraction_gated.parquet" \
MODEL_PATH="$MODEL_PATH" \
OUT_DIR="$EVAL_ROOT/base-model-inject-self-gated/" \
CUDA_VISIBLE_DEVICES="0,1,2,4" \
NUM_GPUS=4 \
GPU_MEM=0.8 \
MICRO_BSZ=2 \
bash examples/hrlib/40_eval.sh

Train GRPO on injected prompts

TRAIN_DATA_PATH="$HOME/data/math/$MODEL_NAME/train_self_abstraction_gated.parquet" \
MODEL_PATH="$MODEL_PATH" \
OUT_DIR="$CHECKPOINT_ROOT/grpo-self-injected" \
CUDA_VISIBLE_DEVICES="0,1,2,4" \
NUM_GPUS=4 \
GPU_MEM=0.8 \
MICRO_BSZ=2 \
bash examples/hrlib/00_grpo.sh

LOCAL_DIR="$CHECKPOINT_ROOT/grpo-self-injected/global_step_58/actor" \
bash examples/test_time_training/merge.sh

MODEL_PATH="$CHECKPOINT_ROOT/grpo-self-injected/global_step_58/merged_hf_model" \
OUT_DIR="$EVAL_ROOT/grpo-self-injected/" \
CUDA_VISIBLE_DEVICES="0,1,2,4" \
NUM_GPUS=4 \
GPU_MEM=0.8 \
MICRO_BSZ=2 \
bash examples/hrlib/40_eval.sh

DATA_VAL="$HOME/data/MATH-500/$MODEL_NAME/test_self_abstraction_gated.parquet" \
MODEL_PATH="$CHECKPOINT_ROOT/grpo-self-injected/global_step_58/merged_hf_model" \
OUT_DIR="$EVAL_ROOT/grpo-self-injected-inject-gated/" \
CUDA_VISIBLE_DEVICES="0,1,2,4" \
NUM_GPUS=4 \
GPU_MEM=0.8 \
MICRO_BSZ=2 \
bash examples/hrlib/40_eval.sh

Compare with teacher extraction

python3 examples/hrlib/evaluate_results.py lift -- \
--baseline "$EVAL_ROOT/base-model/0.jsonl" \
--treated "$EVAL_ROOT/grpo/0.jsonl"

python3 examples/hrlib/evaluate_results.py lift -- \
--baseline "$EVAL_ROOT/base-model-inject-gated/0.jsonl" \
--treated "$EVAL_ROOT/base-model-inject-self-gated/0.jsonl"

python3 examples/hrlib/evaluate_results.py lift -- \
--baseline "$EVAL_ROOT/grpo-injected/0.jsonl" \
--treated "$EVAL_ROOT/grpo-self-injected/0.jsonl"

python3 examples/hrlib/evaluate_results.py lift -- \
--baseline "$EVAL_ROOT/grpo-injected-inject-gated/0.jsonl" \
--treated "$EVAL_ROOT/grpo-self-injected-inject-gated/0.jsonl"

Judge abstraction use

MODEL="deepseek/deepseek-v4-flash" \
FALLBACK_MODEL="deepseek/deepseek-v4-flash" \
VAL_JSONL="$EVAL_ROOT/base-model-inject-gated/0.jsonl" \
OUT_DIR="$EVAL_ROOT/base-model-inject-gated/judge-abs-use" \
bash examples/hrlib/judge_abstraction_use.sh

Test-time Training

The section below preserves the upstream main README. It documents the generic Test-Time Training fork that the HRLib pipeline above builds on — read it for plain GRPO/TTRL/Intuitor training and evaluation without abstraction injection.

Env Setup

git clone https://github.com/moment-timeseries-foundation-model/Test-Time-Training.git

cd Test-Time-Training/verl

conda create -n verl python==3.12
conda activate verl
USE_MEGATRON=0 bash scripts/install_vllm_sglang_mcore.sh
pip install -U "numpy==2.2.0"  # resolve conflicts with Numba
pip install --no-deps -e .

Datasets

To prepare the training dataset, run

python -m examples.data_preprocess.math_dataset_ttt --local_save_dir ~/data/math --data_source DigitalLearningGmbH/MATH-lighteval

To prepare MATH-500 for evaluation, run

python -m examples.data_preprocess.math_dataset_ttt --local_save_dir ~/data/MATH-500 --data_source HuggingFaceH4/MATH-500

The parquet files will be generated under ~/data/.

Train

To train a model, adjust the training config (see "Creating Custom Configs") and run

bash examples/test_time_training/train.sh

The current implementation supports the following reward managers:

reward_manager reward
naive 0/1 based on ground-truth
ttrl majority voting
intuitor self-certainty

Evaluate

To evaluate a trained model, first merge the verl checkpoints into a huggingface model:

bash examples/test_time_training/merge.sh

Then adjust the evaluation config (see "Creating Custom Configs") and run

bash examples/test_time_training/evaluate.sh

Creating Custom Configs

This project uses an extensible config setup, allowing you to override defaults for specific tasks.

To create a custom config, add a YAML file under verl/trainer/config/exps.
Make sure to include # @package _global_ at the top so that overrides work correctly.

To use custom configs, add them to EXPS="[...]" in train.sh or evaluate.sh (the examples/hrlib/*.sh runners declare their own EXPS stacks).
Configs are applied from left to right, with later ones overriding earlier ones. CLI overrides (key=value on the main_ppo command line) beat all exps configs, which is how the scripts' env knobs (MICRO_BSZ, GPU_MEM, N_SAMPLES, ...) keep working.

⚠️ Caution — keep EXPS and naming variables in sync. EXPS is the single source of truth for what actually runs: the advantage estimator and reward manager come from the exps YAMLs. Shell variables like REWARD_MANAGER and ADV_ESTIMATOR in examples/test_time_training/train.sh and examples/hrlib/00_grpo.sh are naming labels only — they build the output directory and the W&B run name, and editing them does not change training behavior. When switching algorithms, change EXPS (e.g. [grpo,intuitor,math]) and update the label variables together, otherwise checkpoints land in a directory named after the wrong algorithm.

Analysis

The code to run analysis regarding prompt accuracy and diversity/entropy for base and trained models can be found in scripts/analyze.py.

About

Code for "Notes to Self: Can LLMs Benefit from Experiential Abstractions?"

Resources

Readme
Activity
Custom properties

Stars

0 stars

Watchers

0 watching

Forks

0 forks
Report repository

Releases

Packages

Contributors

Languages
Morty Proxy This is a proxified and sanitized view of the page, visit original site.