SoftPipe: A Soft-Guided Reinforcement Learning Framework for Automated Data Preparation

TL;DR

SoftPipe introduces a flexible reinforcement learning framework for automated data preparation, replacing rigid constraints with soft guidance from a language model, leading to better and faster pipeline optimization.

Contribution

It proposes a novel RL framework that uses soft guidance via Bayesian inference and LLMs, improving data pipeline quality and convergence speed.

Findings

Up to 13.9% improvement in pipeline quality

2.8× faster convergence than existing methods

Effective on 18 diverse datasets

Abstract

Data preparation is a foundational yet notoriously challenging component of the machine learning lifecycle, characterized by a vast combinatorial search space. While reinforcement learning (RL) offers a promising direction, state-of-the-art methods suffer from a critical limitation: to manage the search space, they rely on rigid ``hard constraints'' that prematurely prune the search space and often preclude optimal solutions. To address this, we introduce SoftPipe, a novel RL framework that replaces these constraints with a flexible ``soft guidance'' paradigm. SoftPipe formulates action selection as a Bayesian inference problem. A high-level strategic prior, generated by a Large Language Model (LLM), probabilistically guides exploration. This prior is combined with empirical estimators from two sources through a collaborative process: a fine-grained quality score from a supervised Learning-to-Rank (LTR) model and a long-term value estimate from the agent's Q-function. Through extensive experiments on 18 diverse datasets, we demonstrate that SoftPipe achieves up to a 13.9\% improvement in pipeline quality and 2.8$\times$ faster convergence compared to existing methods.