P ROMPTS AND P RE -T RAINED L ANGUAGE M ODELS FOR O FFLINE R EINFORCEMENT L EARNING
A BSTRACT
In this preliminary study, we introduce a simple way to leverage pre-trained language models in deep offline RL settings that are not naturally suited for textual representation.
We propose using a state transformation into a human-readable text and a minimal fine-tuning of the pre-trained language model when training with deep offline RL algorithms.
This approach shows consistent performance gains on the NeoRL MuJoCo datasets.
Our experiments suggest that LM finetuning is crucial for good performance on robotics tasks.
However, we also show that it is not necessary when working with finance environments in order to retain significant improvement in the final performance.
I NTRODUCTION
The interest in applying natural language to reinforcement learning is rapidly growing.
One of this sphere’s underexplored directions is the incorporation of pre-trained language models into deep RL algorithms where the environments are not naturally suited for textual representation.
Recently, Reid et al.
(2022) took a step in this direction by taking an earlier proposed Decision
Transformer (Chen et al., 2021) and matching it with a pre-trained language model, demonstrating remarkable results on several deep offline RL benchmarks.
However, this method is designed specifically for the Decision Transformer, and there is a plethora of other algorithms available.
In this work, we demonstrate that there is a simpler way to combine pre-trained language models (PLMs) with deep offline RL algorithms for environments with vector-like observations.
We propose using prompt engineering (Liu et al., 2021a) by presenting the vector state as a string, embedding it via P-LM, and then concatenating the resulting encoding with the original observation.
This results in state-of-the-art performance on a subset of the NeoRL benchmark (Qin et al., 2021) across modelfree algorithms.
B ACKGROUND 2. O FFLINE RL
Reinforcement Learning (RL) is commonly defined by Markov Decision Process (MDP) with a tuple ( S, A, P, r, γ ) , where S is a state space, A is an action space, P ( s ′ | s, a ) is a transition function, r ( s, a ) is a reward function, and γ is a discount factor.
The goal of the agent is to obtain a policy π ( s, a ) that maximizes E π [ (cid:80) ∞ t =
γ t r t +1 ] by interacting with the MDP.
Offline RL is different from RL, as its agent cannot interact with the environment and the policy must be obtained by learning from the dataset D collected with some behavior policy π β or a set of such policies.
The main challenge of offline RL is that D usually does not contain all possible transactions while the agent does not have any opportunity to explore.
This is one of the reasons why this problem is considered challenging, as the agent cannot improve on incorrect estimations of states or actions if they are not presented in D (Levine et al., 2020).
2. L ANGUAGE M ODELS IN RL
There is a considerable amount of recent work that employs pre-trained LMs for solving RL tasks.
For example, Majumdar et al.
(2020) use pre-trained ViLBERT (Lu et al., 2019) to solve vision-andlanguage navigation tasks, and Yao et al.
(2020) show that fine-tuning GPT-
(Radford et al., 2019) for action generation boosts performance in text-based games.
In another work, Singh et al.
(2021) fine-tune DistilBERT (Sanh et al., 2019) to encode states and actions in text-based games, and this approach obtains higher scores comparing to those that do not use pre-trained LMs.
These studies show that knowledge stored in pre-trained LMs might be useful for solving RL tasks with natural language.
For a more comprehensive overview, see Luketina et al.
; Madureira & Schlangen (2020).
While the aforementioned works focus on problems involving natural language, there is a recent interest in incorporating pre-trained LMs for problems not defined via natural language by default.
In particular, Li et al.
(2022) demonstrates that utilizing pre-trained LM for encoding both the goal, history, and parts of the observation for imitation learning leads to improved systematic generalization.
To the best of our knowledge, Reid et al.
(2022) is the only existing work in offline reinforcement learning that adapts a pre-trained LM model for text-free RL tasks.
The authors propose using a pretrained LM to improve the convergence speed and the final performance of the Decision Transformer (Chen et al., 2021) by making the reward, action, and state encodings closer to the language model embeddings.
P ROPOSED A PPROACH
[Figure]
Inspired by the recent success in informing deep RL algorithms with pre-trained language models, we propose an alternative approach to leverage them in environments that are not obviously suited for the textual representation, similar to Reid et al.
; Li et al.
Here, we suggest a method based on prompt-engineering (Liu et al., 2021a) depicted in Figure 1: (1) re-write an observation into a human-readable textual representation (2) encode it using a pre-trained language model, and (3) concatenate the resulting embedding with the original observation.
The first step can be an arduous task, requiring semantic parsing and bringing more structure to the problem.
To alleviate this, we propose using a simple transformation that only requires giving a human-readable description of the dimensions of the state vector.
For example, proprioceptive inputs are common in deep RL literature, and each dimension can usually be mapped to ”Position
X: 10.0 meters” or ”Left Thigh Angle: 0.4 rad”.
So the procedure is as follows: given a vector state s , we apply an environment-specific transformation f : S → String that converts each dimension value into a string, prepends a human-readable label to it, and then concatenates all of these strings into one.
In general, while pre-trained language models seem to handle human-readable texts well, token embeddings provided by pre-trained language models may not transfer between different domains: e.g., the corpus used for training the LM does not contain a lot of text with numbers representing robot joints or velocities.
Therefore, in addition to re-writing the observation, we also perform minimal fine-tuning of the LM when training an offline RL algorithm.
This is done in a way similar Generalizable to Lu et al.
(2021): instead of tuning the entire model, we only tune the input embeddings and layer norm parameters.
The proposed approach can be applied to any off-the-shelf offline RL algorithm, given that the length of the textual representation is suitable for the utilized LM.
Note that it is not necessary to parse the observation space for meaningful entities, as it is enough to simply provide the labels for each dimension.
E XPERIMENTAL E VALUATION
Concrete Realization The proposed method allows for various choices of language models, observation encodings, and deep offline RL algorithms.
We use MPNET (Song et al., 2020) trained with Sentence Transformer (Reimers & Gurevych, 2019) as a backbone for the LM.
For the textual representation, we round the numbers to 5 decimal places to speed up LM inference and reduce the memory space used by the model.
A batch size of 32 is employed in order for the model to fit on a single NVIDIA Tesla V100 GPU with 32 GB of memory.
We rely on TD3+BC (Fujimoto & Gu, 2021) as a deep offline RL algorithm in our experiments.
Benchmarks We probe our approach using a subset of NeoRL benchmark (Qin et al., 2021), running the experiments on the MuJoCo tasks, and report the best performance after a hyperparameter search.
Following the original evaluation protocol by Qin et al.
(2021), we train each policy for 3 · 10 5 iterations.
However, we evaluate the policies with only 100 episodes instead of 1000, as we observed no difference in the final results.
We also run the vanilla TD3+BC algorithm, as its results are not reported in the original NeoRL paper.
[Table]
The results of our preliminary set of experiments are presented in Table 1.
We observe significant performance gains on two out of three settings using the proposed method (TD3+BC, LM-Finetune), and a slight drop on the third one.
Can We Avoid Fine-Tuning?
To assess the importance of fine-tuning the input embeddings, we also run an ablation reported in Table 1 (denoted as LM-No-Tune).
We observe that improvement is possible even without fine-tuning, but only on one of the environments and with a smaller magnitude.
We hypothesize that this is due to a domain shift (Lu et al., 2021), i.e., it is possible that the training data for the LM did not contain much information about robot velocities and angles, and therefore some amount of fine-tuning is helpful.
To test this further, we run the ablation in another environment – FinRL (Liu et al., 2021c).
The observation space is also vector-based, but the values represent different financial indicators (e.g., price of a certain stock, an amount of a specific stock, etc.).
We suppose that these indicators are more common than the proprioceptive inputs when training the
LMs, and therefore a non-tuned LM should result in consistent improvements.
The results for various dataset levels and behavioral policies expertise can be found in Table 6.
And indeed, we can see that no fine-tuning can be sufficient to achieve substantial improvements over the base model in most settings.
[Table]
Sensitivity to Online Evaluation Budget
The preference between deep offline RL algorithms was noted to be dependent on the online evaluation budget available (Kurenkov & Kolesnikov, 2022).
To analyze whether conclusions about the proposed method change under the varied budget, we also report the Expected Online Performance graphs in Figure 2.
Our results indicate that, on environments where an improvement over the base model is possible, said improvement is achieved across all online budgets
[Table]
[Figure]
F UTURE W ORK AND L IMITATIONS
Despite the encouraging results we observed in this preliminary study, we believe there are several important directions that should be explored and experiments to be done to substantiate the claims further:
• Probe on a larger scale : ablations with other deep offline RL algorithms, more environment and dataset sizes.
• Systematic generalization : while we performed testing on a common set of offline RL benchmarks, the most appealing side of the pre-trained language models is their ability to generalize.
And while deep online RL algorithms were studied for systematic generalization (Jiang et al., 2021), their offline counterparts were not.
We expect that the use of language models can substantially benefit deep offline RL algorithms in this aspect.
• Explore different tuning approaches : recent success in prompt-tuning (Liu et al., 2021b) for downstream NLP tasks suggests that the prompt engineering can be bypassed.
We believe that exploring this method type in the context of deep RL can also be of interest.
• Vision-based policies : the proposed algorithm is not well-suited for vision-based policies, and studying how it can be adapted for this problem type is also a potential direction for future research.
A A PPENDIX A. P ROMPTS
[Table]
example of a prompt.
”x position: 0.00394 m, z position: 1.24658 m, y angle: -0.00189 rad, right thigh angle: 0.
rad, right leg angle: -0.00054 rad, right foot angle: -0.00291 rad, left thigh angle: 0.00076 rad, left leg angle: -0.0006 rad, left foot angle: 0.00068 rad, x velocity: 0.00244 m/s, z velocity: -0.
m/s, y angular velocity: -0.00418 rad/s, right thigh angular velocity: 0.00155 rad/s, right leg angular velocity: -0.00413 rad/s, right foot angular velocity: 0.00015 rad/s, left thigh angular velocity: - 0.00351 rad/s, left leg angular velocity: 0.00191 rad/s, left foot angular velocity: -0.00086 rad/s”
”x position: 0.01921 m, z position: 0.05191 m, y angle: -0.06347 rad, back thigh angle: 0.
rad, back shin angle: -0.07455 rad, back foot angle: -0.04487 rad, front thigh angle: 0.01305 rad, front shin angle: -0.05206 rad, front foot angle: 0.04236 rad, x velocity: 0.0732 m/s, z velocity: 0.00765 m/s, y angular velocity: -0.03904 rad/s, back thigh angular velocity: -0.03949 rad/s, back shin angular velocity: -0.01819 rad/s, back foot angular velocity: 0.07899 rad/s, front thigh angular velocity: 0.2711 rad/s, front shin angular velocity: -0.0694 rad/s, front foot angular velocity: 0.02239 rad/s”
[Table]