cs224r-default-project-rloo

cs224r-default-project-rloo

RLOO (REINFORCE Leave-One-Out) fine-tuned model for the Countdown arithmetic reasoning task, built on top of an SFT baseline. Trained as part of Stanford CS224R (Spring 2026).

Model Description

This model is trained with online reinforcement learning using the RLOO algorithm. Given a target number and a set of allowed numbers, the model produces chain-of-thought reasoning inside <think> tags and a final answer inside <answer> tags. A rule-based verifier rewards correct arithmetic equations (score 1.0), correctly formatted but incorrect equations (score 0.1), and malformed outputs (score 0.0).

Training Details

Evaluation

Evaluated on asingh15/countdown_tasks_3to4 test split (50 prompts) using vLLM with temperature 0.6, top-k 20, top-p 0.95, sampling K=16 responses per prompt.

Usage

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("ba144220/cs224r-default-project-rloo")
tokenizer = AutoTokenizer.from_pretrained("ba144220/cs224r-default-project-rloo")

messages = [{"role": "user", "content": "Using the numbers [3, 4, 6, 8], create an equation that equals 24."}]
input_text = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
inputs = tokenizer(input_text, return_tensors="pt").to(model.device)

outputs = model.generate(**inputs, max_new_tokens=1024, temperature=0.6, top_k=20, top_p=0.95, do_sample=True)
print(tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True))

Limitations

• Trained and evaluated only on the Countdown arithmetic task; not intended for general-purpose use.
• Performance degrades on harder problems with more numbers or larger targets.
• The 0.5B parameter size limits reasoning capacity compared to larger models.

Authors

Yuchi Hsu ([email protected]) and Ryan He ([email protected]), Stanford CS224R Spring 2026.