Dhruv Bansal

Graduated from ASU and now building at Blueberry. We're making AI agents that help consumer brands get to know and talk to the people behind their social audience.

Most of what I'm curious about sits around AI agents: how they use context and memory, how they get better, and how to keep them from falling over in production.

I've also written about embedding-based feature selection, multi-agent code generation, and training negotiation agents through self-play.

Before Blueberry, I worked on document-processing systems and data discovery at Chunkr, built internal tools at Oats Overnight, and automated data workflows at PrivateBlok.

I learn best by building small things. I keep a running collection of things I've found interesting—from agent architectures and personal knowledge systems to graphics programming, deep focus, and life design.

Selected work

Projects

Embedding-Based Markov Blanket DiscoveryA causal feature-selection pipeline combining TabPFN embeddings with a neural Markov blanket predictor.

Abstract

Causal feature selection through Markov Blanket (MB) discovery provides theoretical guarantees for optimal prediction while preserving model interpretability. However, identifying the MB without access to the underlying causal graph presents significant challenges, particularly when data exhibits distribution shift across environments. We present a meta-learning framework that combines TabPFN embeddings with neural MB predictors to perform simultaneous feature selection and regression across heterogeneous tasks. Our method trains separate multi-layer perceptron classifiers on TabPFN-extracted representations to predict binary MB masks, which subsequently filter features for downstream regression. Evaluated on 182 meta-training tasks and 46 held-out tasks spanning multiple data-generating processes and environmental conditions, our approach achieves a final score of 0.2071 (RMSE: 0.5078, Jaccard: 0.5921), demonstrating the feasibility of learning MB structure from embeddings without explicit causal discovery while outperforming the baseline by approximately 7.6%.

Introduction

Feature selection is fundamental to machine learning, directly impacting model performance, interpretability, and generalization. Traditional methods often select features based on correlation with the target variable, which can fail under distribution shift when spurious correlations change across environments (Pearl, 1988). Causal feature selection addresses this limitation by identifying features with stable causal relationships to the target.

The Markov Blanket (MB) of a target variable provides a theoretically optimal feature set for prediction. Pearl (Pearl, 1988) showed that the MB renders the target conditionally independent of all other variables, implying that knowing the MB is both sufficient and necessary for optimal prediction. For a target variable YYY, the MB consists of its parents, children, and spouses (parents of children) in the causal graph. This theoretical foundation suggests that identifying the MB should yield robust predictions that generalize across environmental shifts.

Despite this theoretical support, discovering the MB in practice remains challenging. Standard approaches require causal discovery algorithms to first recover the full causal graph, then extract the MB. However, causal discovery methods are computationally expensive, require strong assumptions, and can be unreliable on finite samples (Tsamardinos et al., 2003). This motivates the search for alternative approaches that can identify relevant causal features without explicit graph reconstruction.

We address this challenge by framing MB discovery as a meta-learning problem. Given a distribution over causal data-generating processes, we learn to predict MB masks directly from tabular data. Our approach leverages TabPFN (Hollmann et al., 2023), a transformer-based foundation model for tabular data, to extract rich representations that capture both predictive patterns and structural properties of datasets. These embeddings serve as inputs to learned MB predictors that generalize across tasks with varying feature dimensions and causal structures.

Main Contributions:

  1. A practical pipeline for MB-based feature selection that bypasses explicit causal discovery
  2. Demonstration that TabPFN embeddings encode sufficient information for cross-task MB prediction
  3. Empirical validation across 228 tasks with diverse data-generating processes and environmental conditions

Related Works

Markov Blanket Theory. Pearl (Pearl, 1988) formalized the concept of the Markov Blanket and proved its sufficiency for prediction. Tsamardinos and Aliferis (Tsamardinos et al., 2003) developed algorithms for MB discovery from observational data, establishing principled methods like IAMB for incremental MB learning (Tsamardinos et al., 2022). Recent work has shown that individual causal feature subsets (parents or children alone) can underperform non-causal baselines, emphasizing the importance of the complete MB structure.

TabPFN and Tabular Foundation Models. Hollmann et al. (Hollmann et al., 2023) introduced TabPFN, a transformer trained on synthetic datasets that performs in-context learning for tabular prediction. Unlike traditional models requiring task-specific training, TabPFN processes train-test pairs in a single forward pass, making predictions based on learned priors over data-generating processes. The model's embedding layer provides dense representations that capture both statistical patterns and structural properties of tabular data.

Meta-Learning for Tabular Data. Meta-learning frameworks learn to adapt quickly to new tasks by leveraging experience from related tasks (Finn et al., 2017). In the tabular domain, this typically involves learning good initializations or adaptation strategies. Our approach differs by learning a direct mapping from embeddings to structural properties (MB masks), treating MB discovery as a multi-label classification problem across tasks.

Causal Feature Selection. Traditional causal feature selection methods like the PC algorithm, GES, and NOTEARS first recover the causal graph, then extract relevant features. These methods achieve moderate accuracy with significant computational cost. Our embedding-based approach bypasses explicit graph recovery, directly predicting MB masks in a fraction of the time.

Method & Implementation

Model Design

Our approach consists of a three-stage pipeline that performs MB prediction followed by filtered regression. Each task T\mathcal{T}T is represented by its support set (Xtrain,ytrain)(X_{train}, y_{train})(Xtrain​,ytrain​) and query set (Xtest,ytest)(X_{test}, y_{test})(Xtest​,ytest​) together with metadata such as feature masks, SCM parameters, and environment shifts.

Our pipeline processes tasks independently but uses a shared meta-trained set of models. We handle different feature dimensions by training separate TabPFN regressors, embedding generators, and MBMLP models for each feature size. For example, feature dimensions 19 and 9 are treated as separate groups, each with its own fine-tuned models.

Stage 1: TabPFN Fine-Tuning and Embedding Extraction

The primary technique in our pipeline is batched fine-tuning of the TabPFNRegressor (Prior Labs, 2024). Using the built-in meta dataset collator, each batch corresponds to a full task that includes split preprocessing, differentiable forward passes, and updates to the model weights. This allows the regressor to adapt to the variety of training tasks instead of relying purely on pretraining.

For embedding extraction, we use TabPFNEmbedding to generate fold-based embeddings of each task. For each task with support set (Xtrain,ytrain)(X_{train}, y_{train})(Xtrain​,ytrain​), we employ KKK-fold cross-validation:

  1. Split XtrainX_{train}Xtrain​ into KKK folds
  2. For each fold kkk, train TabPFN on the remaining K−1K-1K−1 folds
  3. Extract embeddings for the held-out fold
  4. Concatenate embeddings across folds

These embeddings average over predictions and internal activations from the regressor, yielding a compact representation of the causal relationships in the data.

Stage 2: MB Prediction via Neural Networks

The meta-training dataset contains tasks with two distinct feature dimensionalities: tasks with 9 features and tasks with 19 features. Rather than using a single unified model with padding, we train separate MB predictors for each dimensionality.

For tasks with d∈{9,19}d \in \{9, 19\}d∈{9,19} features, we define:

MLPd:RE→[0,1]dh1=ReLU(W1x+b1)h2=ReLU(W2h1+b2)m^=σ(W3h2+b3)\begin{aligned} \text{MLP}_d: \mathbb{R}^{E} &\to [0,1]^d \\ h_1 &= \text{ReLU}(W_1 x + b_1) \\ h_2 &= \text{ReLU}(W_2 h_1 + b_2) \\ \hat{m} &= \sigma(W_3 h_2 + b_3) \end{aligned}MLPd​:REh1​h2​m^​→[0,1]d=ReLU(W1​x+b1​)=ReLU(W2​h1​+b2​)=σ(W3​h2​+b3​)​

where EEE is the embedding dimension, and σ\sigmaσ is the sigmoid function. The output m^∈[0,1]d\hat{m} \in [0,1]^dm^∈[0,1]d represents per-feature probabilities of MB membership.

Training Procedure. The MBMLP model performs multi-label binary classification to map embeddings to MB masks. BCEWithLogitsLoss is used during training, and threshold tuning is carried out per feature dimension by evaluating several candidate thresholds on held-out tasks. After training the MBMLP model, we perform threshold tuning for each feature dimension by sweeping thresholds between 0.2 and 0.8, choosing the value that minimizes the combined RMSE and Jaccard-based score.

Inference. At test time, we extract embeddings for a new task, pass them through the appropriate MLPd_dd​, and aggregate predictions via mean pooling followed by thresholding. If no features are selected (all probabilities below threshold), we employ a fallback strategy: select all features to ensure valid predictions.

Stage 3: Filtered Regression

Given predicted MB mask m^∈{0,1}d\hat{m} \in \{0,1\}^dm^∈{0,1}d, we filter both support and query sets:

Xtrain′=Xtrain[:,m^=1]Xtest′=Xtest[:,m^=1]\begin{aligned} X'_{train} &= X_{train}[:, \hat{m} = 1] \\ X'_{test} &= X_{test}[:, \hat{m} = 1] \end{aligned}Xtrain′​Xtest′​​=Xtrain​[:,m^=1]=Xtest​[:,m^=1]​

We then use clone_model_for_evaluation to obtain clean, non-batched TabPFN models. The fine-tuned TabPFN is fitted on the filtered support set (Xtrain′,ytrain)(X'_{train}, y_{train})(Xtrain′​,ytrain​) and used to generate predictions for the filtered query set Xtest′X'_{test}Xtest′​. TabPFN's in-context learning capability allows it to adapt to the filtered feature space without requiring iterative optimization.

Implementation Details

Our entire pipeline is organized in the /solution directory with main.py as the entry point. Data loading from HuggingFace is handled in the BeyondTheBlanket class initialization.

Key files:

  • finetuner.py - TabPFN fine-tuning logic
  • generator.py - Embedding generation
  • mbmlp.py - Markov Blanket MLP predictor

Experimental Setup

Dataset and Task Distribution

We use the CSE472-blanket-challenge/final-dataset from Hugging Face, which contains tasks generated from a hierarchical causal data-generating process. Each task T=(G,SCM,E)\mathcal{T} = (G, \text{SCM}, E)T=(G,SCM,E) is defined by a directed acyclic graph GGG, a structural causal model specifying functional relationships, and an environment E∈{IID,covariate shift,label shift}E \in \{\text{IID}, \text{covariate shift}, \text{label shift}\}E∈{IID,covariate shift,label shift}.

The meta-training set (develop) contains 182 tasks with the following distribution:

  • Feature dimensions: 87 tasks with 9 features, 95 tasks with 19 features
  • Environments: IID, label shift, covariate shift
  • SCM types: Linear and nonlinear
  • Sample sizes: 400 training samples, 100 test samples per task

The held-out set (submit) contains 46 tasks with ground truth MB masks and test labels withheld for final evaluation.

Training Configuration

For meta-training, we grouped tasks by their feature dimensionality. Within each feature dimension group, we held out 10 percent of the tasks for validation and threshold tuning, and used the remaining 90 percent for fine-tuning the TabPFN regressor and training the MBMLP network. Hyperparameter configurations for all pipeline components are provided in Table in the Appendix.

Computational Resources

All experiments were conducted on ASU Sol HPC cluster using A100 GPUs (40GB VRAM). Total pipeline runtime was approximately 20 minutes including fine-tuning, embedding extraction, MBMLP training, and threshold tuning. Submit prediction uses cloned evaluation-mode TabPFN models and is comparatively lightweight.

Results

Main Quantitative Results

The results were calculated by running baseline tests and the final model on held-out tasks (~10% from the develop set).

The baseline delivers an average RMSE of 0.4878, average Jaccard of 0.5504, and combined score of 0.2193. Our pipeline delivers an average RMSE of 0.5078, average Jaccard of 0.5921, and an aggregate final score of 0.2071.

DimModelRMSEJaccardScore
19Baseline0.50710.50880.2491
19Final0.52580.56250.2301
9Baseline0.46610.59720.1878
9Final0.48760.62550.1826

Performance comparison by feature dimension

The final model produces a higher Jaccard score at the cost of slightly elevated RMSE. The final model achieves approximately 5.6% overall improvement over the baseline (7.6% for dimension 19, 2.8% for dimension 9).

Qualitative Analysis

The graph's complexity with entangled paths and dispersed blanket nodes makes MB prediction a particularly challenging problem. Beyond individual examples, a broader inspection of tasks reveals that MB prediction quality is closely tied to the structural clarity of the underlying causal graph. Tasks with well-separated parent and child relationships tend to produce higher Jaccard values, even when the regression RMSE remains moderate.

Figure shows the causal graph for task data_e963d7d6, which achieved high Jaccard score (0.9) despite moderate RMSE (0.7790). The causal graph structure is relatively clear, with well-separated relationships contributing to accurate MB prediction. Figure shows the causal graph for task data_be3555b2, which achieved lower Jaccard score (0.3) with RMSE 0.8202, due to more complex graph structure with entangled paths.

Research figure Causal graph for data_e963d7d6 (RMSE: 0.7790, Jaccard: 0.9)

Research figure Causal graph for data_be3555b2 (RMSE: 0.8202, Jaccard: 0.3)

Analysis & Insights

What Worked and Why

Task-structured training: For each supported feature dimension, develop tasks are split into a "support" subset for fine-tuning TabPFN and training the MB network, and a "query" subset for threshold tuning and evaluation, mirroring the benchmark's support/query structure.

Causal/MB signal: TabPFN embeddings for each task feed an MB-specific MLP that predicts Markov Blanket membership per feature dimension. Averaging per-task MB probabilities and thresholding keeps only features most likely in the causal Markov Blanket before running the regression fine-tune.

Joint score tuning: Held-out tasks drive threshold selection by minimizing a combined objective (RMSE×(1−Jaccard)\text{RMSE} \times (1 - \text{Jaccard})RMSE×(1−Jaccard)), directly balancing predictive accuracy with MB fidelity and exploiting causal sparsity assumptions.

Failure Modes and Limitations

Sparse tasks per dimension: When few develop tasks exist for a feature dimension, the support split and MB training become data-starved, making the MB classifier and threshold tuning noisy.

MB instability: Thresholding can yield empty blankets, triggering the fallback to all features. Frequent fallback indicates the MB model cannot distinguish relevant features—often on high-dimensional or highly correlated tasks.

Distribution shift: The MB network and thresholds are tuned on develop held-outs; submit tasks with substantially different causal structure may see degraded blanket accuracy and downstream RMSE because the learned thresholds no longer align.

Future Work

Architectural Improvements

Unified Multi-Task Model. Our current approach trains separate MLP predictors for each feature dimension. A unified architecture using adaptive pooling or padding could leverage shared structure across dimensions while maintaining specialized capacity. Graph neural networks such as DAG-GNN (Yu et al., 2019) could model feature dependencies explicitly, potentially improving MB prediction for features with complex interdependencies.

Attention-Based Aggregation. We currently aggregate sample-level predictions via mean pooling. An attention mechanism could learn to weight samples based on their informativeness, potentially improving robustness to outliers or low-quality training examples.

Training Enhancements

TabPFN Fine-Tuning Extensions. TabPFN 2.5 supports fine-tuning via batched training mode. Rather than using zero-shot predictions, deeper fine-tuning on develop tasks could adapt the model to the specific distribution of causal data-generating processes, improving both embedding quality and downstream regression performance.

Uncertainty-Aware Selection. Our current threshold of 0.5 for MB membership is fixed. A learned or adaptive threshold based on prediction confidence could improve precision-recall trade-offs. Incorporating established causal discovery ideas like IAMB (Tsamardinos et al., 2022) would also provide a more principled signal for selecting parent and spouse variables.

Evaluation Extensions

Per-Environment Analysis. Our validation metric averages across all environments. Stratified evaluation would reveal whether performance degrades differentially under distribution shift, informing environment-specific adaptations.

Sample Efficiency Studies. Investigating performance as a function of support set size would characterize the method's data efficiency. TabPFN's in-context learning may enable reasonable performance with very few support examples.

Appendix

Hyperparameters

ComponentParameterValue
TabPFN (fine-tuning)epochs1
TabPFN (fine-tuning)learning_rate1×10−61 \times 10^{-6}1×10−6
TabPFN (fine-tuning)max_samples10,000
TabPFN (fine-tuning)val_fraction0.1
MLP Predictorhidden_dims[256, 128]
MLP Predictorlearning_rate1×10−31 \times 10^{-3}1×10−3
MLP Predictorepochs40
MLP Predictorbatch_size512
MLP Predictorval_fraction0.2
TabPFN (regression)n_estimators24

Hyperparameter configuration for all pipeline components

Explicit (Exposed through config.yaml)

HyperparameterDescription
supported_feature_dimsFeature dimensions (e.g., 9, 19)
holdout_fractionFraction held out for evaluation

General Configuration

HyperparameterDescription
tabpfn.epochsFine-tuning epochs
tabpfn.lrLearning rate (Adam)
tabpfn.max_samplesMax samples per dataset
tabpfn.val_fractionInternal validation fraction
tabpfn.model_pathPretrained model checkpoint

TabPFN Fine-tuning Configuration

HyperparameterDescription
mb_mlp.epochsTraining epochs
mb_mlp.lrLearning rate
mb_mlp.batch_sizeTraining batch size
mb_mlp.hidden_sizesHidden layer widths
mb_mlp.val_fractionValidation fraction
mb_mlp.n_foldFolds for embedding extraction

MBMLP Configuration

HyperparameterDescription
threshold_tuning.minMin threshold in sweep
threshold_tuning.maxMax threshold in sweep
threshold_tuning.num_thresholdsGrid search values

Threshold Tuning Configuration

Implicit Hyperparameters

ComponentHyperparameterReason
TabPFNn_estimators = 24Ensemble size
MBMLPactivation = ReLUExpressiveness
MBMLPloss = BCEWithLogitsLossTraining objective
MBMLPoptimizer = AdamLearning trajectory
MBMLPbetas = (0.9, 0.999)Momentum
Thresholdrule = >=Mask selection
Thresholdaggregation = meanMB selection
Thresholdmetric = RMSE×\times×(1-Jacc)Best threshold

Implicit Hyperparameters

Training Loss on MBMLP

Training loss curves for MBMLP models during meta-training:

Research figure MBMLP training loss (9 features)

Research figure MBMLP training loss (19 features)

References

  1. Judea Pearl. Probabilistic Reasoning in Intelligent Systems: Networks of Plausible Inference. Morgan Kaufmann Publishers. 1988.
  2. Ioannis Tsamardinos, Constantin F. Aliferis. Towards Principled Feature Selection: Relevancy, Filters and Wrappers. Proceedings of the Ninth International Workshop on Artificial Intelligence and Statistics. 2003.
  3. Noah Hollmann, Samuel Müller, Katharina Eggensperger, Frank Hutter. TabPFN: A Transformer That Solves Small Tabular Classification Problems in a Second. arXiv preprint arXiv:2207.01848. 2023.
  4. Huan Liu, Hiroshi Motoda. Computational Methods of Feature Selection. Chapman,Hall/CRC. 2007.
  5. Brady Neal. Introduction to Causal Inference. 2020.
  6. AutoML. Prior-Data Fitted Networks (PFNs). 2023.
  7. Prior Labs. Prior Labs Documentation. 2024.
  8. Sebastian Raschka. Feature Selection Categories Overview. 2023.
  9. Chelsea Finn, Pieter Abbeel, Sergey Levine. Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks. Proceedings of the 34th International Conference on Machine Learning. 2017.
  10. Ioannis Tsamardinos, Constantin Aliferis, Alexander Statnikov. The Incremental Association Markov Blanket Algorithm. Algorithms. 2022.
  11. Yue Yu, Jianfei Gao, Song-Chun Zhu. DAG-GNN: DAG Structure Learning with Graph Neural Networks. Proceedings of the 36th International Conference on Machine Learning. 2019.
Advancing Multi-Agent Reinforcement Learning for Collaborative Coding AgentsTwo coding agents trained together with shared rewards and multi-agent group relative policy optimization.

Abstract

Complex programming tasks challenge single-agent LLM systems since they require deep understanding, planning, and iterative debugging. Most existing multi-agent frameworks depend on static prompting or manually designed pipelines, struggling with dynamic communication, credit assignment, and scalability.

We develop a multi-agent framework driven by reinforcement learning for cooperative coding agents. A Centralized Training Decentralized Execution (CTDE) approach enables multiple LLMs to collaborate in efficiently planning and executing coding tasks.

Our aim is to improve the traditional LLM's and Agent implementation using the complexity of Reinforcement Learning in a controlled environment.

Introduction

Large Language Models (LLMs) demonstrate strong capabilities in natural language reasoning and code synthesis, but single-agent systems often struggle with complex programming tasks requiring decomposition, multi-step planning, iterative debugging, or feedback integration. Multi-agent LLM systems have recently emerged as a promising direction, enabling role specialization, communication, and collaborative decision-making (Wu et al., 2023; Qian et al., 2023; Islam et al., 2024).

However, most systems, such as MapCoder (Islam et al., 2024), ChatDev (Qian et al., 2023), and AutoGen (Wu et al., 2023), depend heavily on manually engineered pipelines rather than learning-based coordination. Reinforcement Learning (RL), especially in multi-agent settings, has not been widely applied to collaborative coding, with only limited progress from frameworks like Co-Learning (Yu et al., 2025), MAGRPO (Liu et al., 2025), and TF-GRPO (Cai et al., 2025).

We fill this gap by designing a reinforcement-learning-based multi-agent framework that enables LLM agents to coordinate adaptively using CTDE, collaborative decision-making, and shared rewards driven by test-case outcomes. We enhance the MARL architecture by employing two agents (helper and main) and using Multi-Agent Group Relative Policy Optimization (MAGRPO) that trains multiple LLMs jointly through centralized group-relative advantages while maintaining decentralized execution efficiency. Modeling LLM collaboration as a cooperative MARL problem and training with MAGRPO yields adaptive, sample-efficient systems capable of solving challenging programming tasks through learned coordination strategies rather than hand-engineered pipelines.

Related Work

Multi-Agent Collaboration

Multi-agent LLM frameworks demonstrate that structured roles and communications can improve code generation and debugging. AutoGen (Wu et al., 2023) enables multi-agent conversation but relies on prompt engineering rather than adaptive learning. MapCoder (Islam et al., 2024) uses specialized agents for retrieval, planning, coding, and debugging, achieving strong benchmark results but following a fixed pipeline. ChatDev (Qian et al., 2023) models a full software development workflow using designer, coder, tester, and documenter agents, but lacks dynamic coordination. Co-Learning (Yu et al., 2025) introduces test-driven reinforcement to choose agents based on code correctness, though it depends on handcrafted rewards and a small dataset.

Communication and role specialization are effective, but none of these frameworks support learning-based scheduling, credit assignment, or scalable adaptive coordination.

Reinforcement Learning in Multi-Agent LLM Systems

Few studies apply reinforcement learning to collaborative LLM coding. Co-Learning (Yu et al., 2025) uses Environmental RL to reward correct outputs but does not train agents jointly. Training-Free GRPO (Cai et al., 2025) offers lightweight relative policy updates without pretraining, reducing computation but remaining limited to synthetic benchmarks.

These approaches show early success but do not yet scale to heterogeneous coding agents, motivating the need for a CTDE-driven multi-agent framework.

Data

Our project uses the HumanEval dataset to support code generation evaluation, multi-agent coordination, and reinforcement learning for debugging.

HumanEval

Source: OpenAI HumanEval Benchmark Purpose: Evaluating Functional correctness of generated Python code using unit tests Size: 164 programming tasks and test cases

Why We Use It: HumanEval (Chen et al., 2021) serves as the standard benchmark for assessing LLM coding performance. It allows direct comparison with established multi-agent systems such as MapCoder (Islam et al., 2024) and MAGRPO (Liu et al., 2025).

Dataset Split: HumanEval is typically evaluated as a fixed test set. We follow a standard practice:

  • 75.76% train (25/33 ratio)
  • 24.24% test (8/33 ratio)

Methods

Our approach combines the multi-agent architecture from Co-Learning with the Multi-Agent Group Relative Policy Optimization (MAGRPO) algorithm for end-to-end reinforcement learning. We conduct a systematic evaluation progressing from baseline single-agent systems to fully coordinated multi-agent systems with reinforcement learning.

Our system consists of two specialized agents (main and helper) that work together to solve programming problems. Each agent operates with partial knowledge of the complete system state, observing different aspects of the problem-solving process based on its role.

Multi-Agent Architecture

Our system instantiates two specialized LLM agents, each with distinct roles in the collaborative coding workflow:

Helper Agent (Agent 0): Generates utility/helper
functions that decompose the problem into manageable sub- components. The helper agent receives the problem descrip- tion and entry point, then produces auxiliary functions that the main agent can leverage.

Main Agent (Agent 1): Receives both the original prob- lem prompt and the helper code generated by Agent 0. It then completes the target function, potentially utilizing the helper functions to construct the final solution.

MAGRPO Algorithm

We adapt Group Relative Policy Optimization (GRPO) to multi-agent settings as described in (Liu et al., 2025). At each training episode:

  1. Sample a problem and initialize state
  2. For each turn t∈[0,H−1]t \in [0, H-1]t∈[0,H−1]: a. Generate GGG diverse solutions per agent b. Execute actions and compute rewards c. Store trajectory data
  3. Compute group-relative advantages: At(g)=Rt(g)−mean(Rt)A_t^{(g)} = R_t^{(g)} - \text{mean}(R_t)At(g)​=Rt(g)​−mean(Rt​)
  4. Update policies using policy gradient with loss: L=−A⋅log⁡πθ(a∣h)L = -A \cdot \log \pi_\theta(a|h)L=−A⋅logπθ​(a∣h)

We use Qwen2.5-Coder-3B as our base model. Our implementation differs from (Liu et al., 2025) in several ways: they use Qwen2.5-Coder-7B (larger model), a tree-like structure for rollouts, discounted returns with configurable discount (γ=0.9\gamma = 0.9γ=0.9), and multiple optimizer steps per tree node.

Experiments

Hardware: 1x NVIDIA A100 (40GB or 80GB), 2× AMD EPYC 7413 processors (24 cores each), 503 GiB RAM total, Rocky Linux 8.10 (Green Obsidian) Base Model: Qwen2.5-Coder-3B (bfloat16) Dataset: HumanEval

Evaluation Protocol

  1. Baseline Single Model: The base Qwen2.5-Coder-3B model without any fine-tuning. For each problem, we generate k independent solutions using temperature sampling (temperature=0.8) via the vLLM inference server. Each solution is generated from the problem prompt alone, with code extracted from the model's response. This establishes our baseline Pass@k performance for a single untrained agent.

  2. Baseline Multi-Agent: Two instances of the same Qwen2.5-Coder-3B model working in sequence without RL training. Agent 0 (helper) receives a specialized prompt asking it to generate a helper function for the problem. Agent 1 (main) then receives both the original problem and the helper's generated code, completing the target function. This tests whether role decomposition through prompting alone improves performance over the single-agent baseline.

  3. GRPO Single-Agent: A single Qwen2.5-Coder-3B model trained using Group Relative Policy Optimization. The model generates G candidate solutions per problem, receives rewards based on code correctness (+0.5 for runnable code, +0.1 per passed test up to 5 tests), and updates using group-relative advantages. Training runs for multiple epochs over the training split, with the final model evaluated on the test set.

  4. Multi-Agent GRPO: Our full system with two specialized agents (helper and main) trained jointly using MAGRPO. Both agents share the same reward signal based on the combined code's correctness. The helper agent learns to generate useful auxiliary functions while the main agent learns to leverage them effectively. Rewards are computed on the combined output: +0.2 for syntactically valid code, +0.3 for runnable code, and +0.1 per passed test. Agents are updated using centralized group-relative advantages while maintaining decentralized execution.

Metrics

Pass@k: Probability that at least 1 of kkk solutions passes all tests:

pass@k=Eproblems[1−(n−ck)(nk)]\text{pass@}k = \mathbb{E}_{\text{problems}} \left[ 1 - \frac{\displaystyle \binom{n-c}{k}}{\displaystyle \binom{n}{k}} \right]pass@k=Eproblems​​1−(kn​)(kn−c​)​​

where nnn is the total number of samples, ccc is the number of correct samples, and kkk is the number of samples considered.

Results

Method@1@3@5@10
Single Model29.25%53.02%63.11%72.50%
Multi-Agent24.00%49.50%61.94%75.00%
GRPO14.50%30.96%40.96%55.00%
MAGRPO44.25%65.23%72.29%80.00%

Results on 40/164 HumanEval problems.

Discussion

Our results reveal several key insights:

  1. Multi-Agent MAGRPO achieves the best performance across all Pass@k metrics, with a 51% relative improvement in Pass@1 over the single-agent baseline (44.25% vs 29.25%).

  2. Baseline multi-agent underperforms the single-agent baseline at lower k values (24% vs 29.25% at Pass@1), suggesting that naive role decomposition through prompting alone can introduce coordination overhead.

  3. Single-agent GRPO shows degraded performance, likely due to small model size and limited training data, which may cause mode collapse. This could also be attributed towards small group size (G) of only 4.

  4. MAGRPO’s joint training is important. The improvement from baseline multi-agent to MAGRPO shows that learned coordination significantly outperforms static prompting-based collaboration.

Conclusion

We presented a multi-agent reinforcement learning framework for collaborative code generation using MAGRPO. Our results demonstrate that jointly training specialized agents (helper and main) with shared rewards substantially outperforms both single-agent baselines and prompt-based multi-agent systems.

The key finding is that learned coordination through MAGRPO enables effective collaboration that static role assignment cannot achieve. While baseline multi-agent systems may struggle with coordination overhead, RL training allows agents to develop complementary strategies that leverage each other's outputs effectively.

Future work includes scaling to larger models, exploring more complex agent hierarchies, and extending evaluation to additional benchmarks.

References

  1. Qingyun Wu, others. AutoGen: Enabling Next-Gen LLM Applications via Multi-Agent Conversation. arXiv preprint arXiv:2308.08155. 2023.
  2. Chen Qian, others. ChatDev: Communicative Agents for Software Development. arXiv preprint arXiv:2307.07924. 2023.
  3. Md Ashraful Islam, others. MapCoder: Multi-Agent Code Generation for Competitive Problem Solving. Proceedings of the 62nd Annual Meeting of the ACL. 2024.
  4. Jiaying Yu, others. Co-Learning: Code Learning for Multi-Agent Reinforcement Collaborative Framework. arXiv preprint arXiv:2409.00985. 2025.
  5. Shengguang Liu, others. LLM Collaboration with Multi-Agent Reinforcement Learning. arXiv preprint arXiv:2508.04652. 2025.
  6. Yifan Cai, others. Training-Free Group Relative Policy Optimization. arXiv preprint arXiv:2510.08191. 2025.
  7. Mark Chen, others. Evaluating Large Language Models Trained on Code. arXiv preprint arXiv:2107.03374. 2021.
  8. Shuai Lu, others. CodeXGLUE: A Machine Learning Benchmark Dataset for Code Understanding and Generation. NeurIPS Datasets and Benchmarks Track. 2021.
Improving LLM Bargaining Ability via Self-Play Reinforcement LearningTraining open-weight LLM negotiators with supervised fine-tuning, self-play reinforcement learning, and outcome-based rewards.

Abstract

Automated negotiation with large language models (LLMs) is commercially attractive but technically demanding: a capable agent must track private utilities, adapt strategy across many turns, produce numerically valid proposals, and know when to concede or walk away. We consider two regimes in which LLMs participate in negotiation. In the mixed setting, an LLM faces a human or unstructured counterpart in an assistive or advisory role. In the symmetric multi-agent setting, two LLM agents bargain under a shared formal protocol with private priorities until a deal, impasse, or turn limit is reached. This paper presents a complete pipeline for the symmetric regime: a structured Thought / Talk / Action output format that decouples strategic intent from language generation, a three-stage training procedure combining supervised fine-tuning with reinforcement learning, and a composite reward scheme targeting deal-level outcomes. We train and evaluate on the CaSiNo multi-issue bargaining corpus, pitting the trained agent in head-to-head play against opponents conditioned on diverse personas. Our results show that the trained agent modestly improves aggregate deal rate relative to the untuned base model, while also highlighting that agreement rate and utility are distinct objectives that can move independently. We discuss the implications of these findings for the design and evaluation of negotiation agents in business settings.

Introduction

Negotiation is a routine but high-stakes function across procurement, commercial contracting, business development, and resource allocation. Automating it with AI is a compelling opportunity, yet the task resists the kind of single-turn, single-prompt reasoning at which modern LLMs excel. A capable negotiating agent must maintain private utility estimates across many turns, adapt its strategy as the counterpart reveals preferences through offers and rejections, and ultimately produce numerically valid agreements that reflect both tactical strength and situational judgment.

We distinguish two common regimes in which LLMs participate in negotiation. In the mixed setting, an LLM occupies one side of the table while the counterpart is a human or a separately-prompted model without a shared protocol. In the symmetric multi-agent setting, two LLM agents negotiate under a shared formal protocol, each observing only the public dialogue and their own private priorities, until the episode closes with a deal, a walk-away, or a turn limit.

The central challenge motivating this work is that linguistic fluency and negotiation competence are not the same thing. A systematic evaluation spanning 35 tasks across four negotiation corpora shows that even frontier models produce contextually coherent language while making irrational offers, conceding utility too readily, and failing to model their counterpart's priorities through the conversation (Kwon et al., 2024). Models can score well on comprehension and annotation tasks while remaining poor real-time negotiators.

This paper presents our approach to building a locally trainable, open-weight negotiation agent for the symmetric multi-agent setting.

Read the anonymous preprint

A running collection of things I find interesting

Wiki

AI Agents & Agentic Engineering

Context EngineeringBuilding AI AgentsAgentic Engineering Best PracticesAuto-Research and Self-Improving AgentsSelf-Maintaining AgentsScaling Long-Running Autonomous CodingData AgentsAgent Memory and Harness Ownership

LLM Internals & Training

Reinforcement Learning for LLMsCursor Composer 2Training-Free GRPOPrompt Repetition for Non-Reasoning LLMsLLM QuantizationLLM Knowledge Bases

AI Safety & Security

Constitutional ClassifiersProject Glasswing

Graphics & Low-Level Programming

Shaders and GPU Programming

Personal Development & Mindset

Deep Focus and ProductivityIdentity and Life DesignSelf-Competition and Entrepreneurship