Policy Learning with a Natural Language Action Space: A Causal Approach

TL;DR

This paper presents a causal, data-efficient reinforcement learning framework for natural language action spaces, enabling effective policy learning with limited data and translating embeddings into coherent language.

Contribution

It introduces a single-model Q-learning approach with a decoding strategy for natural language actions, improving data efficiency and transfer strength in language-based decision tasks.

Findings

Outperforms baselines in mental health, hate speech, and sentiment tasks

Achieves better transfer strength while preserving content and fluency

Demonstrates practical policy learning in limited-data scenarios

Abstract

This paper introduces a novel causal framework for multi-stage decision-making in natural language action spaces where outcomes are only observed after a sequence of actions. While recent approaches like Proximal Policy Optimization (PPO) can handle such delayed-reward settings in high-dimensional action spaces, they typically require multiple models (policy, value, and reward) and substantial training data. Our approach employs Q-learning to estimate Dynamic Treatment Regimes (DTR) through a single model, enabling data-efficient policy learning via gradient ascent on language embeddings. A key technical contribution of our approach is a decoding strategy that translates optimized embeddings back into coherent natural language. We evaluate our approach on mental health intervention, hate speech countering, and sentiment transfer tasks, demonstrating significant improvements over competitive baselines across multiple metrics. Notably, our method achieves superior transfer strength while maintaining content preservation and fluency, as validated through human evaluation. Our work provides a practical foundation for learning optimal policies in complex language tasks where training data is limited.