Source-linked AI summary
COGTRL: Training LLMs for Scientific Discovery Assistance using Cognitive Traces via Reinforcement Learning
Shrinidhi Kumbhar Santosh Mashetty Divij Handa Kevin Coutinho, Siddharth Sambhaji Ghule, Chitta Baral
TL;DR
Scientific discovery assistance requires methods that pursue goals under constraints, yet prior work largely applies frontier models to narrower scientific tasks. COGTRL trains open-source LLMs with cognitively grounded trace–step trajectories, improving method generation across domains and matching larger models.
Problem
Scientific discovery assistance requires generating step-by-step methodologies for goals under constraints, while prior work primarily applies frontier LLMs to narrower scientific tasks.
Method
COGTRL uses trajectory-level reinforcement learning to jointly generate cognitive traces and corresponding scientific steps, with rewards for trace and method quality.
Results
COGTRL improves method quality by an average of 7.85 points over comparable 3B baselines, matches 70B+ models, and receives expert preference 71.42% of the time.
Takeaways & Limitations
Cognitive traces are most beneficial when they causally improve downstream methodological quality rather than serving as standalone explanations.
Takeaways & Limitations
COGTRL relies on closed-source language models for reward calculation and rubric-based rewards rather than physical-world, tool-based, or simulation-grounded verification.
Abstract
from arXiv · showhide
Large Language Models (LLMs) trained on extensive scientific research are increasingly integrated as assistants for scientific discovery. However, most research papers omit the fine-grained cognitive process of examining constraints, failed alternatives, and iterative decisions required to achieve the desired goal. Such cognitive processes are vital for real-world scientists working toward specific goals under constraints. In this paper, we show that LLMs, when trained to produce such cognitive traces, perform better as scientific discovery assistants than when trained solely on scientific literature. We propose COGTRL, a trajectory-level reinforcement learning framework that trains LLMs to emulate cognitively grounded reasoning by jointly optimizing cognitive traces and the scientific steps produced in an interleaved manner. Across two 3B-parameter models and two scientific domains (AI and Materials Science), COGTRL improves method quality by an average of 7.85 points over comparable 3B model baselines and achieves competitive performance relative to 70B parameter models. Moreover, analysis by domain experts shows a preference for methods generated by COGTRL over the baselines.
1 Introduction
COGTRL trains open-source LLMs to generate cognitively grounded traces alongside constraint-aware scientific methodologies. Across two 3B models and two domains, it improves method generation and receives expert preference over baselines.
- Motivation: Scientific discovery assistance asks LLMs to generate step-by-step methodologies that achieve research goals while satisfying explicit constraints.Such assistance may reduce the time and labor costs of methodological design.
- Motivation: Scientific literature often omits intermediate cognitive processes such as evaluating constraints and iteratively refining plans.The paper argues that training should instill these behaviors because scientific discovery is fundamentally cognitive.
- COGTRL: COGTRL uses trajectory-level reinforcement learning to jointly train cognitive traces and scientific steps for open-ended discovery assistance.Its trajectories generate traces before each method step, with later steps conditioned on preceding traces and steps.
- Results: 71.42% of the time, domain experts preferred COGTRL-generated methods in blind evaluation across AI and Materials Science.The 3B COGTRL models also matched the performance of much larger 70B+ models and preserved or modestly improved out-of-domain reasoning.
2 Related Work
Prior scientific-discovery work uses agentic workflows and training methods for specialized tasks, while reasoning-trace methods mainly target verifiable domains. COGTRL combines cognitively grounded traces with reinforcement learning for open-ended, constraint-aware method generation.
- Scientific-discovery systems: Recent scientific-discovery systems use retrieval-augmented, literature-centric, and multi-agent workflows for question answering, synthesis, hypothesis generation, and experimental planning.These approaches leverage citations, constraints, graph reasoning, grounded exploration, and autonomous execution.
- Training approaches: Training-based scientific-discovery approaches include domain-adaptive pretraining, supervised fine-tuning, reinforcement learning, rationale distillation, and process supervision.They support tasks including scientific question answering, information extraction, material generation, molecular design, and physical-system modeling.
- Reasoning traces: Reasoning-trace methods primarily target mathematics, question answering, or programmatically verifiable tasks.They include rationale bootstrapping, process supervision, reinforcement-learning-based reasoning, and latent reasoning.
- COGTRL: COGTRL integrates cognitively grounded intermediate traces with reinforcement learning for open-ended, constraint-aware scientific method generation.It jointly rewards cognitive traces and downstream methodological step quality.
3 Method
COGTRL models scientific method generation as interleaved cognitive traces and methodological steps, then jointly rewards trace quality, step quality, and valid trajectory structure during GRPO training.
- Trajectory formulation: COGTRL generates trajectories by interleaving cognitive traces with corresponding methodological steps conditioned on previous traces and steps.The input consists of a research goal and constraints, and trajectories are generated autoregressively.
- Reward design: The trajectory-level reward evaluates both cognitive traces and methodological steps using scores produced by a large reasoning model.Trace and step scores are aggregated for policy optimization, with an additional structural reward for valid formatting.
- Reward design: Rtrace scores traces from 1 to 5 on dimensions including goal and constraint integration, scientific reasoning, causal logic, and actionability.The reward targets reasoning useful for scientific method generation.
- Reward design: Rstep scores generated methodological steps from 1 to 5 for alignment with objectives and constraints, scientific plausibility, novelty, and testability.The reward evaluates the scientific quality of the generated method itself.
- Uplift reward: Ruplift = Rstep · σ(Rtrace − α) reinforces traces that improve the quality of their conditioned methodological steps.The experiments set α = 0.6, and total reward combines step, uplift, and structural rewards as Rtotal = Rstep + γRuplift + λRstruct.
- Policy optimization: GRPO updates a joint trace–step policy using group-normalized advantages computed from total trajectory rewards.COGTRL samples groups of trajectories, aggregates trace and step rewards into Rtotal, and applies the clipped GRPO objective with a KL stability constraint.
4 Experimental Setup
The experiments compare COGTRL with multiple baselines using two 3B open-source models, two generation formats, and AI and Materials Science evaluation sets. Methodologies are assessed with independent agentic evaluators, while additional safeguards address evaluation reliability and reward hacking.
- Baselines and models: COGTRL is compared against zero-shot Chain-of-Thought, SFT without traces, SFT with traces, and vanilla GRPO without traces.Training experiments use Llama-3.2-3B-Instruct and Qwen-2.5-3B-Instruct; zero-shot comparisons also include larger and frontier-scale models.
- Training data: The training dataset is constructed from arXiv papers spanning six domains using agents that extract goals, constraints, and step-by-step methodologies.The extracted data achieves an average correctness score of 2.58/3 and Krippendorff’s alpha of 0.71 in annotation validation.
- Evaluation benchmark: The held-out benchmark covers 50 NeurIPS 2024 AI papers and 50 Materials Science papers from the Mat-Design benchmark, with expert-extracted methodological steps.Additional out-of-domain reasoning benchmarks measure performance beyond scientific methodology generation.
- Evaluation: Generated methodological steps are evaluated independently by Codex- and Claude Code-based agentic frameworks with search and tool use enabled.Evaluators are distinct from the o3-mini reward model used during reinforcement learning, reducing reward-model overfitting; results are averaged across three runs.
- Evaluation: Table 2 reports quality results for COGTRL and baselines across the Think First and Int. Think generation formats.The evaluation focuses on generated methodological steps rather than cognitive traces.
5 Results and Analysis
COGTRL improves scientific method quality across models, domains, baselines, and evaluation settings. Evidence indicates that interleaved, aligned cognitive traces and uplift rewards contribute to these gains while preserving general reasoning capability.
- Overall performance: 7.85 points: COGTRL improves over non-COGTRL baselines using the same 3B models across AI and Materials Science.It improves over vanilla GRPO by an average of 4.12 points in the interleaved setting.
- Overall performance: 3B COGTRL models become competitive with open-source 70B models, outperforming Qwen-2.5-72B-Instruct on AI and matching or exceeding Llama-3.3-70B-Instruct.
- Generation strategy: 4.30 points on AI and 4.99 on Material Science: interleaved COGTRL outperforms think-first COGTRL across both models.Similar trends for SFT with traces suggest that conditioning each step on preceding traces and steps yields stronger alignment than generating reasoning upfront.
- Reward design: 5.38%: removing the uplift reward causes average degradation across both models, indicating traces help when they improve downstream method quality.COGTRL also shows similar gains with Dr. GRPO, suggesting it is not tied to a specific RL algorithm.
- Trace evidence: Aligned traces yield higher probabilities for fixed high-quality methodological steps than no-trace or mismatched-trace settings across models and domains.The comparison controls for target step tokens, supporting trace–step alignment rather than context length or formatting as the source of gains.
- Step-quality dimensions: +4.55: COGTRL’s largest average per-dimension gains are in Testability and Impact Potential, followed by +4.50 in Scientific Plausibility.Alignment with Constraints and Goals, Innovation and Novelty, and Feasibility and Scalability each improve by +3.65.
- Out-of-domain performance: COGTRL preserves average general benchmark performance across both model families while improving scientific method generation.It avoids the degradation seen with SFT across several benchmarks and has modest gains over vanilla GRPO on AMC23 and GPQA-Diamond.
- Human evaluation: 71.42%: domain experts prefer COGTRL methods across AI and Materials Science, highlighting completeness, organization, and causal structure.In Materials Science, experts specifically report greater specificity, actionability, constraint integration, and systematic workflows.
6 Conclusion
The paper presents COGTRL as trajectory-level reinforcement learning for jointly generating cognitive traces and scientific methodologies. Across experiments and human evaluations, it improves scientific method generation and remains competitive with larger models while preserving general reasoning capabilities.
- Contribution: COGTRL jointly generates cognitive traces and scientific methodologies through trajectory-level reinforcement learning with multidimensional rewards.
- Conclusion: 3B-parameter COGTRL models achieve competitive performance relative to larger 70B+ models while preserving general reasoning capabilities.
- Conclusion: Cognitive traces are most beneficial when they causally improve downstream methodological quality rather than serving as standalone explanations.
Limitations
The study identifies limitations in reward computation, verification, evaluation bias, and model scope. It also notes future directions beyond the current study.
- Reward calculation relies on closed-source LRMs such as OpenAI o3-mini, introducing additional API cost.
- The framework uses rubric-based LLM rewards rather than physical-world, tool-based, or simulation-grounded verification.
- LLM-as-a-judge evaluation may inherit biases and may not perfectly capture real-world scientific validity.
Ethics Statement
The paper reports using AI assistants to correct grammar and rephrase sentences. It specifically names Grammarly and ChatGPT.
- The authors used AI assistants for grammatical correction and sentence rephrasing.
- The disclosed assistants were Grammarly and ChatGPT.
- The paper separately reports that COGTRL outperforms several baselines by an average of 7.85 points.
B SFT Data Construction
The SFT data pipeline converts research papers into structured examples containing goals, constraints, methodologies, and cognitive traces. It combines automated extraction, generated annotations, critic selection, and methodological prompts.
- Dataset construction: The dataset samples 2K arXiv paper titles from each of six scientific domains and parses papers into structured JSON.
- Information extraction: Agents extract research goals, constraints, and step-by-step methodologies from the structured paper representations.
- Cognitive-trace annotation: GPT-4o generates three cognitive-trace candidates per step at temperatures 0.3, 0.5, and 0.7.
- Cognitive-trace annotation: A multi-agent critic ensemble scores candidate thoughts on a 1–5 scale before selection.
- Scoring: Evaluation aggregates six quality dimensions by dividing their summed scores by 30 and multiplying by 100.
- Methodology prompting: Prompts require detailed, sequential methodologies that include actions, conditions, resources, parameters, and technical specifications.
G Training Details
The training setup uses distinct configurations for supervised fine-tuning and reinforcement learning, with rollouts generated through vLLM. Training curves track reward trajectories for both models and algorithms.
- SFT: SFT uses batch size 4, learning rate 2.0 × 10^-5, cosine scheduling, 5 epochs, and AdamW on 2 NVIDIA H100 GPUs.
- Reinforcement learning: GRPO and COGTRL use learning rate 1 × 10^-6, mini-batch size 12, and micro-batch size 2 per GPU.
- Reinforcement learning: GRPO trains for 400 steps and COGTRL for 200 steps, with vLLM used for rollouts.
- Training analysis: Training figures report cumulative reward trajectories for GRPO and COGTRL across Llama-3.2 3B-Instruct and Qwen-2.5 3B-Instruct.
H.1 Material Science
The Materials Science example evaluates methodologies for designing a sustainable photothermal-responsive self-healing elastomer under coupled mechanical, optical, and healing constraints. COGTRL’s method links synthesis, characterization, and modeling decisions causally to those constraints, and experts rank it above the alternatives.
- Three coupled constraints require high strength with repeatable healing, balanced photothermal conversion and agent dispersion, and fast localized healing under NIR irradiation.Increasing agent loading can cause aggregation and mechanical degradation, while increasing crosslink density can slow healing kinetics.
- Method C: Method C jointly designs photothermal materials and varies agent concentration and incorporation strategy to target dispersion-controlled NIR absorption rather than maximum heating alone.It also analyzes morphology-dependent NIR penetration depth and temperature uniformity.
- Method C: Method C designs hierarchical crosslinked networks and validates healing under cyclic loading and fatigue so healing remains compatible with sustained mechanical stress.
- Method C: Method C uses temperature monitoring, agent-size and geometry optimization, and iterative characterization to support fast localized repair and refine a healing-kinetics model.The stated characterization includes in-situ FTIR, XPS, and interfacial imaging.
- The target is a sustainable photothermal-responsive self-healing elastomer for smart anticorrosion coatings and flexible electronics.
- Expert evaluation: Experts ranked COGTRL first, zero-shot second, and GRPO third, describing COGTRL as novel, realistic, and supported by logical characterization and analysis stages.
I.2.1 Expert Evaluation
Expert evaluation compares COGTRL with zero-shot and GRPO methodologies for realistic LiDAR generation. COGTRL is preferred because it covers more stated constraints, while the reward design ties cognitive-trace quality to successful methodological steps.
- Expert evaluation: COGTRL systematically targets sparse processing, physics-based dynamic simulation, domain adaptation, efficiency, robustness, sensor diversity, and metric-based evaluation.
- Expert evaluation: GRPO provides a clear hybrid neural-plus-reconstruction framework and dynamic-scene simulator, but is less explicit about robustness, sensor diversity, and aggressive efficiency optimization.
- Expert evaluation: Zero-shot uses a standard reconstruction pipeline but treats simulation, domain gap, sparsity, and efficiency as add-ons rather than core design drivers.It therefore only partially satisfies the stated constraints and lacks systematic targeting of dynamic scenes, domain gap, and strong efficiency guarantees.
- Reward formulation: The reward formulation includes Rstep as a dense stabilizing reward and inside Ruplift as a causal tether linking trace-quality bonuses to successful actions.This coupling rewards reasoning most when its correctness is substantiated by a high-quality causal reasoning trace.
M Rollout examples of CoGTRL model training
The rollout examples illustrate how COGTRL rewards interleaved cognitive traces and methodological steps that integrate constraints, mechanisms, and actionable decisions. Higher-reward examples contain stronger mechanistic detail and cleaner trace-step structure than lower-reward examples.
- Rollout examples: A lower-quality rollout receives reward 0.20 after duplicate or malformed tags, very short steps, absent mechanism, and failure to address signal transmission.
- Rollout examples: An intermediate rollout receives reward 0.51 for clean tag structure and emerging mechanistic detail, but only partial constraint integration and incomplete treatment of signal transmission.
- Rollout examples: Its trace connects reversible bonding, hydrophobic microdomains, and conductive networks to rapid healing, limited swelling, and pressure sensing in aquatic conditions.
- Rollout examples: A high-quality rollout receives reward 0.67 for proper tags, strong constraint integration, explicit mechanistic reasoning, preserved signal transmission, and scalable processing.
- Extraction prompts: The associated extraction prompts require concise goals, constraint-focused problem definitions, and detailed, accurate, terminology-preserving methodology steps.
- Trace generation: The cognitive-trace generation procedure requires forward-looking causal reasoning before each methodological step, using only information in the research text and avoiding solution leakage.
O Test Time Evaluation Judge Settings
Test-time evaluation uses Codex and Claude Code as LLM judges with specified model, decoding, and web-search settings.
- Codex and Claude Code serve as LLM judges, using GPT-5.4 and Claude Sonnet 4.6, respectively.
- Both judges evaluate only generated methodological steps, use default temperature 1 and maximum-token settings, and have web search enabled.
P Independent Rubric Evaluation
COGTRL is evaluated with an independent rubric using Claude Sonnet 4 and Claude Opus 4.8, and achieves the highest overall performance across both model families and judges. The evaluation pipeline scores generated steps through independent judgments and applies group-normalized GRPO-style updates during training.
- Independent rubric results: The independent evaluation assesses Consistency, Clarity, Novelty, Feasibility, Significance, and Overall score.These criteria come from the MLR-Bench Idea Generation evaluation.
- Evaluation design: COGTRL is evaluated with an independent rubric that was not used during training, supporting generalization beyond the training evaluation setup.The evaluation uses the MLR-Bench rubric as an additional safeguard against reward hacking.
- Independent rubric results: COGTRL achieves the highest overall performance for both model families across both independent judges.Table 8 uses Claude Sonnet 4 and Claude Opus 4.8 as judges.
- Training procedure: Training samples multiple trajectories per prompt, parses each into trace–step pairs, and computes separate step and trace judgments.The resulting rewards include step quality, trace quality, uplift, and structural components before policy updating.
- Training procedure: Group-normalized advantages are computed per prompt before a clipped GRPO-style policy update with KL regularization.The policy ratio compares the current policy with the previous policy, while the reference policy supplies KL regularization.