Pretraining Decision Transformers with Reward Prediction for In-Context Multi-task Structured Bandit Learning

• URL: http://arxiv.org/abs/2406.05064v1
• Date: Fri, 7 Jun 2024 16:34:31 GMT
• Title: Pretraining Decision Transformers with Reward Prediction for In-Context Multi-task Structured Bandit Learning
• Authors: Subhojyoti Mukherjee, Josiah P. Hanna, Qiaomin Xie, Robert Nowak,
• Abstract summary: We study multi-task structured bandit problem where the goal is to learn a near-optimal algorithm that minimizes cumulative regret. We use a transformer as a decision-making algorithm to learn this shared structure so as to generalize to the test task. We show that our algorithm, without the knowledge of the underlying problem structure, can learn a near-optimal policy in-context.
• Score: 12.608461657195367
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: In this paper, we study multi-task structured bandit problem where the goal is to learn a near-optimal algorithm that minimizes cumulative regret. The tasks share a common structure and the algorithm exploits the shared structure to minimize the cumulative regret for an unseen but related test task. We use a transformer as a decision-making algorithm to learn this shared structure so as to generalize to the test task. The prior work of pretrained decision transformers like DPT requires access to the optimal action during training which may be hard in several scenarios. Diverging from these works, our learning algorithm does not need the knowledge of optimal action per task during training but predicts a reward vector for each of the actions using only the observed offline data from the diverse training tasks. Finally, during inference time, it selects action using the reward predictions employing various exploration strategies in-context for an unseen test task. Our model outperforms other SOTA methods like DPT, and Algorithmic Distillation over a series of experiments on several structured bandit problems (linear, bilinear, latent, non-linear). Interestingly, we show that our algorithm, without the knowledge of the underlying problem structure, can learn a near-optimal policy in-context by leveraging the shared structure across diverse tasks. We further extend the field of pre-trained decision transformers by showing that they can leverage unseen tasks with new actions and still learn the underlying latent structure to derive a near-optimal policy. We validate this over several experiments to show that our proposed solution is very general and has wide applications to potentially emergent online and offline strategies at test time. Finally, we theoretically analyze the performance of our algorithm and obtain generalization bounds in the in-context multi-task learning setting.

Related papers

• Data-CUBE: Data Curriculum for Instruction-based Sentence Representation Learning [85.66907881270785]
  We propose a data curriculum method, namely Data-CUBE, that arranges the orders of all the multi-task data for training. In the task level, we aim to find the optimal task order to minimize the total cross-task interference risk. In the instance level, we measure the difficulty of all instances per task, then divide them into the easy-to-difficult mini-batches for training.
  arXiv  Detail & Related papers  (2024-01-07T18:12:20Z)
• Supervised Pretraining Can Learn In-Context Reinforcement Learning [96.62869749926415]
  In this paper, we study the in-context learning capabilities of transformers in decision-making problems. We introduce and study Decision-Pretrained Transformer (DPT), a supervised pretraining method where the transformer predicts an optimal action. We find that the pretrained transformer can be used to solve a range of RL problems in-context, exhibiting both exploration online and conservatism offline.
  arXiv  Detail & Related papers  (2023-06-26T17:58:50Z)
• Hierarchically Structured Task-Agnostic Continual Learning [0.0]
  We take a task-agnostic view of continual learning and develop a hierarchical information-theoretic optimality principle. We propose a neural network layer, called the Mixture-of-Variational-Experts layer, that alleviates forgetting by creating a set of information processing paths. Our approach can operate in a task-agnostic way, i.e., it does not require task-specific knowledge, as is the case with many existing continual learning algorithms.
  arXiv  Detail & Related papers  (2022-11-14T19:53:15Z)
• Fast Inference and Transfer of Compositional Task Structures for Few-shot Task Generalization [101.72755769194677]
  We formulate it as a few-shot reinforcement learning problem where a task is characterized by a subtask graph. Our multi-task subtask graph inferencer (MTSGI) first infers the common high-level task structure in terms of the subtask graph from the training tasks. Our experiment results on 2D grid-world and complex web navigation domains show that the proposed method can learn and leverage the common underlying structure of the tasks for faster adaptation to the unseen tasks.
  arXiv  Detail & Related papers  (2022-05-25T10:44:25Z)
• In Defense of the Unitary Scalarization for Deep Multi-Task Learning [121.76421174107463]
  We present a theoretical analysis suggesting that many specialized multi-tasks can be interpreted as forms of regularization. We show that, when coupled with standard regularization and stabilization techniques, unitary scalarization matches or improves upon the performance of complex multitasks.
  arXiv  Detail & Related papers  (2022-01-11T18:44:17Z)
• Learning to Actively Learn: A Robust Approach [22.75298609290053]
  This work proposes a procedure for designing algorithms for adaptive data collection tasks like active learning and pure-exploration multi-armed bandits. Our adaptive algorithm is learned via adversarial training over equivalence classes of problems derived from information theoretic lower bounds. We perform synthetic experiments to justify the stability and effectiveness of the training procedure, and then evaluate the method on tasks derived from real data.
  arXiv  Detail & Related papers  (2020-10-29T06:48:22Z)
• Meta-learning with Stochastic Linear Bandits [120.43000970418939]
  We consider a class of bandit algorithms that implement a regularized version of the well-known OFUL algorithm, where the regularization is a square euclidean distance to a bias vector. We show both theoretically and experimentally, that when the number of tasks grows and the variance of the task-distribution is small, our strategies have a significant advantage over learning the tasks in isolation.
  arXiv  Detail & Related papers  (2020-05-18T08:41:39Z)
• Learning Adaptive Exploration Strategies in Dynamic Environments Through Informed Policy Regularization [100.72335252255989]
  We study the problem of learning exploration-exploitation strategies that effectively adapt to dynamic environments. We propose a novel algorithm that regularizes the training of an RNN-based policy using informed policies trained to maximize the reward in each task.
  arXiv  Detail & Related papers  (2020-05-06T16:14:48Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.