Towards Execution-Grounded Automated AI Research — part 3/4

Cleaned-text code points [31065, 47021). Bibliography omitted here; source.md includes it.

Recent works have been building LLM-based research agents for accelerating scientific discovery in various domains, including AI research. AI-Scientist ( Lu et al., 2024b ; Yamada et al., 2025 ) , AI-Researcher ( Tang et al., 2025 ) , and Agent Laboratory ( Schmidgall et al., 2025 ) are examples of end-to-end research agents that use LLMs to generate ideas and implement them through carefully designed agent scaffolds. They address open-ended AI research as we do, but do not study how to learn from execution feedback and improve idea effectiveness. On the other hand, on more grounded benchmarks with clear performance metrics such as MLE-Bench ( Chan et al., 2025 ) , RE-Bench ( Wijk et al., 2024 ) , and ML-Gym ( Nathani et al., 2025 ) , various works have explored how to learn from execution feedback through either search ( Toledo et al., 2025 ; Jiang et al., 2025 ) or RL ( Yang et al., 2025b ) to optimize performance on these targeted ML engineering tasks. While we also study algorithms for learning from execution feedback, we tackle open-ended research problems like pre-training and post-training rather than ML engineering tasks that heavily depend on feature engineering and hyper-parameter tuning rather than algorithm development.

#### AI for Research

Apart from fully end-to-end automated AI research, many works have studied how to use LLMs for specific components of the scientific research pipeline, such as literature review ( Asai et al., 2024 ; L’ala et al., 2023 ) , idea generation ( Si et al., 2025b ; Wang et al., 2024 ) , data analysis ( Majumder et al., 2025 ; Mitchener et al., 2025 ) , experiment plan generation ( Goel et al., 2025 ) , research code execution ( Starace et al., 2025 ; Hua et al., 2025 ; Tian et al., 2024 ) , and paper reviewing ( Liang et al., 2024 ; Zhu et al., 2025 ) . Our work focuses on automated idea execution and learning from the execution feedback. We consider our work complementary to many of the above works that improve other aspects of the scientific research pipeline.

#### Execution Grounding for Code

The idea of learning from execution feedback has been explored in the code generation domain. For example, Zheng et al. (2024) curate data and train models to refine code from either human or execution feedback; Gehring et al. (2025) use end-to-end RL training to teach models to improve code based on execution feedback; Lavon et al. (2025) directly guide code generation with execution signals during inference time. In contrast, our work explores execution grounding for the application of idea generation, where the verification is more complicated and expensive.

## 7 Conclusion

In this work, we built a large-scale parallel executor for automatically executing model-generated ideas to verify their effectiveness on open-ended LLM research problems, including both LLM pre-training and post-training. Using this executor as a reward function, we analyzed the effectiveness of execution-guided evolutionary search, where frontier LLMs equipped with a simple evolutionary search scaffold can significantly outperform the baseline solutions. We also benchmarked and revealed the limitations of reinforcement learning with execution rewards, where models tend to converge on simple ideas to improve the average reward but lose diversity and do not improve the upper-bound. Our empirical results demonstrate the feasibility and potential of the automated execution feedback loop and also point out the remaining limitations for future improvement. We hope this work paves the foundation and inspires more efforts towards execution-grounded automated AI research.

## 8 Discussion

Despite encouraging initial signals, there are still many limitations to our current set of experiments, which would be great directions for future improvement.

First, our current procedure does not test the generalizability of the generated ideas. It is possible that the best-performing ideas at the small scales may not transfer to gains at a larger scale or on other datasets. Future work should explore methods that explicitly test such generalizability and scalability, and even incorporate them as part of the optimization objectives.

Second, we have shown that RL with execution reward in our current setup can only improve the average reward but not the upper bound. There are many possible reasons, such as a lack of diversity in the base model and missing exploration incentives in the current RL objective. Future work should explore remedies and better learning algorithms for LLMs to more efficiently learn from the execution feedback. For instance, future works could explore how to exploit richer learning signals from the execution trajectories beyond just scalar rewards.

Third, our current experiment scope is bounded by the capability of the execution agent. There exist many promising model-generated ideas that could not be successfully executed by the execution agent (e.g., see the end of Appendix A.2 ), leading to noise in the reward signal. Future work could develop more capable execution agents and extend our setup to even more complex research problems. For instance, instead of directly prompting an LLM for code diff, future work can implement more capable coding agents with access to external tools and the ability to install new libraries in the execution environments.

Last but not least, we only explored effectiveness as the training reward in this work. There are many other, more subjective alternative metrics that could complement the effectiveness reward, such as the idea novelty and interestingness. Future work could explore how to computationally measure them and incorporate them as part of the training objective to discover more insightful ideas.

## Acknowledgment

We thank DST Global, Laude Institute, and Thinking Machines Lab for their generous sponsorship of computing resources. We thank Yuandong Tian, Edward Hughes, Ludwig Schmidt, Cong Lu, Jenny Zhang, Jiaxin Wen, Chris Rytting, Chen Zhao, Xinran Zhao, Yanzhe Zhang, John Yang, Shicheng Liu, Andy Zhou, Will Held, Haotian Ye, Luke Bailey, as well as members of Tatsu Lab and SALT Lab for their helpful discussion and feedback. This work was supported by an HAI grant, DSO labs, gifts from Open Philanthropy, Amazon, Schmidt Sciences, the Tianqiao and Chrissy Chen Foundation and a grant under the NSF CAREER IIS-2338866 and IIS-2247357, ONR N00014-24-1-2609 and N00014-24-1-2532, and DARPA Cooperative Agreement HR00112520013. This work does not necessarily reflect the position or policy of the government and no official endorsement should be inferred.

## Appendix A Appendix

### A.1 Other RL Attempts

We present several attempts to improve our RL from the execution reward recipe.

#### Attempt 1: Dynamic Prompt

At each epoch (except the first epoch), we randomly sample different executed idea trajectories from the previous epoch and append them to the idea sampling prompt when sampling new rollouts. This merges in-context learning with RL and adds diversity to the idea sampling process. We present the experiment results on the GRPO environment in Figure 8 . We did not see significant improvement in early epochs and thus early stopped.

#### Attempt 2: Length Reward

Since we noted a rapid thinking length decrease in our main RL experiment, we tried a simple fix by adding a weighted length reward that counts the number of tokens in the entire rollout sequence, including the thinking trace and the idea. We cap the length reward to a maximum of 0.3 0.3 to avoid it dominating the accuracy reward. We present the experiment results on the GRPO environment in Figure 9 . As shown on the right panel, the thinking length no longer decreases after adding the length reward to the total reward, but the total training reward isn’t going up as shown on the left panel.

#### Attempt 3: Diversity Reward

We also tried adding a diversity reward in addition to the effectiveness reward. Specifically, when computing the reward for each rollout, we compute its token-level Jaccard similarity with ideas from the previous epoch and add the negative similarity as a penalty to the total reward to discourage repeating ideas that have already been generated before. In fact, this is similar to one of the post-training ideas proposed by Claude-4.5-Sonnet (see example 4 in Appendix A.3 ). We show the training curves on the GRPO environment in Figure 10 . The model maintains a consistent idea similarity (right plot), suggesting sustained exploration. The effectiveness reward is generally showing an upward trend (left plot), but not markedly better than the main RL run with just the effectiveness reward (first sub-plot in Figure 5 ).

### A.2 Additional Idea Examples

We provide several additional example ideas generated by Claude-4.5-Opus (Table 4 ) and Claude-4.5-Sonnet (Table 5 ) on the GRPO environment, including ideas with failed code execution. In most cases, code execution errors happen when the idea involves complicated changes or installing and importing external packages not supported in our execution environment. Future work should explore improvement to the execution agent so that more complicated types of ideas (e.g., training additional auxiliary models or system-level optimizations) can be implemented correctly.

We also present the top-performing ideas from Claude-4.5-Opus, Claude-4.5-Sonnet, and GPT-5 on the nanoGPT environment below:

While the best-performing ideas on nanoGPT tend to be heavily optimized with extensive hyper-parameter tuning mixed with various architecture tweaks, we also pick a few more “atomic” algorithmic ideas that are successfully executed.

Examples from Claude-4.5-Opus

• Head-Wise Attention Output Scaling Add learnable per-head scaling factors to attention, allowing different heads to contribute with different magnitudes to the output. Validation Loss: 3.2386

• Learned Residual Connection Weights Add learnable scalar weights for each residual connection that are initialized to 1.0, allowing the model to learn optimal residual scaling during training. Validation Loss: 3.2517

• Mixture of Embeddings with Position Learn to mix token embeddings and position embeddings with a content-dependent weight, allowing the model to dynamically balance positional vs semantic information per token. Validation Loss: 3.2497

• Shared Input-Output Embedding with Learned Asymmetry Keep weight tying but add a small learned transformation on the output side, providing the benefits of weight tying while allowing output-specific adaptation. Validation Loss: 3.2499

• Gated Final Normalization Replace the final RMSNorm before lm_head with a gated version where a learned gate controls how much normalization is applied vs passing the raw representation. Validation Loss: 3.2503

• Position-Aware MLP Gating Gate the MLP output based on position information, allowing the model to learn position-dependent processing depth. Validation Loss: 3.2506

• Grouped Token Embeddings Group the vocabulary into clusters and add a learned embedding per cluster on top of token embeddings, providing hierarchical vocabulary structure. Validation Loss: 3.2521

Lastly, we present several interesting ideas on the nanoGPT environment that didn’t get successfully executed. These examples are generated by Claude-4.5-Opus.

• Soft Layer Repetition Allow the model to softly repeat computation through layers by adding a learned gate that mixes the current layer’s input back into its output, simulating variable depth.

• Causal Context Compression Before each attention layer, apply a learned compression that mixes local context (previous 2-3 tokens) into the current representation, providing implicit local context without convolutions.

• Attention Head Specialization via Orthogonal Loss Add a soft penalty that encourages different attention heads to attend to different patterns by penalizing similarity between head outputs.

• Skip Connections with Learned Residual Weights Combine skip connections with learned residual weights. The skip connections provide alternative gradient paths while learned weights allow adaptive scaling.

• Token Difficulty-Aware Loss Weighting Weight the loss contribution of each token based on the model’s uncertainty (entropy) at that position, focusing learning on difficult tokens while not over-optimizing easy ones.

### A.3 Code Execution Examples

We present a few ideas with their full code execution to demonstrate the full end-to-end trajectories. All examples below are from Claude-4.5-Sonnet on the GRPO environment. For each example, we first present the natural language idea, followed by the code implementation generated by Claude-4.5-Sonnet.

Example 1

[Experiment] Create mathematical working memory simulation by maintaining a context buffer of mathematical facts, definitions, and intermediate results during problem solving. This buffer gets updated as the model works through problems and provides additional context for subsequent mathematical steps, simulating how humans maintain mathematical working memory during complex calculations.

[Code Changes] Modify the prompt template in sample.py to include a “Mathematical Context” section that accumulates relevant mathematical facts during problem solving. Implement context buffer management in train_loop in grpo.py that maintains recent mathematical definitions, formulas, and intermediate calculations across problems in the same batch. Update the buffer with key mathematical expressions extracted from each response. Use enhanced prompts that include relevant context from the buffer. Set --learning_rate 2.5e-5 and --loss_type reinforce_with_baseline .

Validation Accuracy: 58.0

⬇

Example 2

[Experiment] Implement mathematical solution robustness training through systematic perturbation testing that teaches the model to solve mathematically equivalent problems with varied presentations, notation styles, and problem phrasings. Combine this with proven reward shaping by creating robustness-aware rewards that encourage mathematical understanding that generalizes across problem variations.

[Code Changes] Create problem perturbation system in sample.py that generates equivalent mathematical problems with varied notation, different variable names, alternative problem phrasing, and equivalent mathematical formulations. Track solution consistency across perturbations. Enhance r1_zero_reward_fn_train in drgrpo_grader.py to reward robustness: robustness_bonus = 0.12 * consistency_score_across_perturbations where consistency measures solution method stability across equivalent problems. Apply generalization bonuses: +0.08 for solving problems with unfamiliar notation, +0.06 for handling varied problem presentations. Modify train_loop in grpo.py to include perturbation-based training batches and track robustness improvement. Use robustness-aware learning rate with cosine annealing from 3.4e-5 to 1.8e-5 . Implement perturbation-complexity gradient clipping. Set --loss_type reinforce_with_baseline .

Validation Accuracy: 56.2

⬇

Example 3

[Experiment] Implement token-level reward attribution by using attention weights to identify which input tokens contributed most to correct answers, then amplifying the gradient updates for those tokens during policy gradient training.

[Code Changes] Modify get_response_log_probs in utils.py to also return attention weights from the last layer. In grpo_microbatch_train_step , compute token importance scores by averaging attention weights across heads, then multiply the policy gradient loss by (1 + importance_score) for tokens with high attention to mathematical content.

Validation Accuracy: 45.2

⬇

Example 4