Table of Contents
Fetching ...

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

Bohan Zhang, Yixin Wang, Paramveer S. Dhillon

TL;DR

This work addresses learning optimal policies in multi-stage tasks with natural language actions by casting policy learning as a causal problem and solving it with offline Q-learning that operates on language embeddings. A single-model framework optimizes a Q-function through gradient ascent on text embeddings and decodes the resulting embeddings into fluent natural language actions, enabling data-efficient learning without explicit text-space search. The NLPolicyLearn approach, including a novel embedding-based decoding strategy, is evaluated on mental health intervention refinement, hate speech countering, and sentiment style transfer, showing improved transfer strength and balanced fluency and content preservation compared to strong baselines, with human evaluations corroborating the gains. The results demonstrate the practicality of causal, embedding-based policy learning for complex language-generation tasks under limited training data, offering a foundation for scalable, multi-stage NL decision-making systems.

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.

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

TL;DR

This work addresses learning optimal policies in multi-stage tasks with natural language actions by casting policy learning as a causal problem and solving it with offline Q-learning that operates on language embeddings. A single-model framework optimizes a Q-function through gradient ascent on text embeddings and decodes the resulting embeddings into fluent natural language actions, enabling data-efficient learning without explicit text-space search. The NLPolicyLearn approach, including a novel embedding-based decoding strategy, is evaluated on mental health intervention refinement, hate speech countering, and sentiment style transfer, showing improved transfer strength and balanced fluency and content preservation compared to strong baselines, with human evaluations corroborating the gains. The results demonstrate the practicality of causal, embedding-based policy learning for complex language-generation tasks under limited training data, offering a foundation for scalable, multi-stage NL decision-making systems.

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.

Paper Structure

This paper contains 33 sections, 4 equations, 3 figures, 6 tables, 1 algorithm.

Figures (3)

  • Figure 1: Multi-stage decision-making with a natural language action space. For simplicity, the figure shows a two-stage example. The task is to transfer a sequence of textual interventions for mental health issues from ineffective to effective. The texts with ineffective signals are highlighted in red and the optimal natural language policy given by the LM are highlighted in blue. Steps within the same stage of Q-Learning are represented using the same colored arrows.
  • Figure 2: The maximization process of Q-learning in each stage under our framework.
  • Figure 3: The loss curve of PPO training on the DIALOCONAN datasets. The algorithm cannot converge with limited training data and delayed rewards.