[ICLR 2022] Part 3: Reinforcement Learning as a sequence modeling problem

This year marks the 10th anniversary of the International Conference on Learning Representations (ICLR), the most prominent conference in the world in the field of deep learning that was first held in 2013. It is also a venue for the dissemination of cutting-edge research and publications in all fields that employ deep learning, including artificial intelligence, data science, vision, speech recognition, and robotics. We introduce some of the papers presented by LG AI Research at the online ICLR 2022 conference held in April. - Part 1: Efficient Transformer — Junghee Kim, Applied AI Research Lab - Part 2: Time Series Anomaly Detection — Yeseul Shim, Applied AI Research Lab - Part 3: Reinforcement Learning as a sequence modeling problem — Deunsol Yoon, Applied AI Research Lab

Reinforcement learning is a method for solving problems involving continuous decision-making. It is a method employed for discovering an agent's behavioral policy that enables the agent to perceive the current state of the environment and choose the action that maximizes the cumulative reward. The agent interacts with the environment and executes multiple actions in order to determine the optimal policy by obtaining reward feedback on each action. Figure 1 is a schematic representation of general reinforcement learning.

(Figure 1) General Reinforcement Learning Structure

Under the name “deep reinforcement learning,” reinforcement learning combined with deep learning has recently demonstrated satisfactory performance in a variety of problems. However, despite the fact that reinforcement learning is a highly promising technology, it is not utilized as frequently as vision or natural language processing (NLP). One of the reasons for this is the lack of a simulator. In deep reinforcement learning, a large amount of data is required to train an agent, but it is difficult to acquire data due to the high risk associated with directly executing various actions in the real world. Consequently, a simulator that can easily collect data is required, but the majority of the real-world problems we wish to solve lack a suitable simulator.

Offline reinforcement learning (Offline RL) has recently been in the spotlight as a solution to this problem. In contrast to conventional reinforcement learning, offline RL is a method that directly trains an agent through given data without involving the environment. Although it is difficult for agents to directly act and obtain data in the real world, there is a vast amount of data created by people. If we can directly train an agent with this data, it will be much easier to apply in the real world. Figure 2 is a schematic representation of offline RL.

(Figure 2) Offline Reinforcement Learning Structure

In the field of offline RL, a new paradigm called Decision Transformer (DT) was proposed in June of last year. In DT, the reinforcement learning problem was solved by transforming it into a problem commonly encountered in NLP — sequence modeling. In particular, DT garnered attention in the field of reinforcement learning due to its successful application of the transformer structure to reinforcement learning and training of the reinforcement learning agent via supervised learning. In this blog, we introduce DT, a recently popularized method of training reinforcement learning via supervised learning, as well as Generalized DT, a model that extends this.

Decision Transformer: Reinforcement Learning via Sequence Modeling

In conventional reinforcement learning, the agent's policy is learned using an action-value function. Action-value function Q(s,a) is the expected value of the sum of rewards that can be earned by acting a in state s. In other words, it can be seen as an expression of how valuable the action is in the current state. And we can obtain the policy π(s) from the action-value function, which instructs on what action to take in state s.

In contrast to vision or NLP, where dozens of neural network layers are stacked, deep reinforcement learning typically expresses action-value functions and policies by stacking only two or three neural network layers. The reason for this is that it is extremely challenging to stably train an action-value function. Temporal Difference (TD) learning, which is utilized when training the action-value function, is one of the primary causes of instability in the learning of the action-value function. The TD learning equation is shown below.

In TD learning, when training the action-value function for the current state s and action a, the action-value function for the next state s’ is used to create the target. As the unstable action-value function currently being trained is being applied to the target, this results in unstable learning. Therefore, in conventional deep reinforcement learning, which uses TD learning to train the action-value function, the use of a layer with a complex structure has been avoided to prevent learning instability.

What is DT?

In this study, the reinforcement learning problem is viewed as a sequence modeling problem, and a new method for training policy via supervised learning is proposed. As DT^[1] does not use TD learning, it could learn a stable policy while employing a layer with a complex structure, such as a transformer. The operation of the DT is depicted in Figure 3.

(Figure 3) Decision Transformer Operation Process Proposed in this Study

In conventional deep reinforcement learning, an action-value function is learned to measure the value of each action, and a policy is learned based on this to determine what action to take in a given situation; however, DT has no action-value function. Instead, it immediately learns what to do when the desired value is to be extracted from the pertinent state. And to train this, it uses transformer, an autoregressive model.

 

How to train?

The training method for the DT closely resembles the training method for the seq2seq model in NLP. Given a trajectory $\left({\mathrm{s}}_1,a_1,r_1,\dots ,s_T\mathit{,}{a}_T\mathit{,}{r}_T\right)$ consisting of a sequence of state, action, and reward, the return-to-go $\overbrace{{\mathrm{R}}_{\mathrm{t}}}=\sum _{t\text{'}=t}^T{r}_{t\text{'}}$, the cumulative reward from each time point t to the end of the episode, is computed. The trajectory is then reconstructed as $\widehat{{\mathrm{R}}_1},s_1,a_1,\dots ,\widehat{{\mathrm{R}}_T},s_T,a_T$ by adding return-to-go and trained using a subsequence with the same length as K. In other words, given the input $\left(\widehat{{\mathrm{R}}_{t-K-1}},s_{t-K-1},a_{t-K-1},\dots ,\widehat{{\mathrm{R}}_t},s_t\right)$, the transformer is trained to predict a_t. If the behavior is discrete, cross entropy loss is used, and if the behavior is continuous, mean square error loss is used. Below is the code for training DT.

In general reinforcement learning, the policy is derived from the action-value function Q(s,a) whose output is return-to-go. However in DT put return-to-go goes into the policy as an additional input by utilizing the conditional policy. As such, the policy learns immediately that it must take action a in order to obtain return-to-go in the current state s.

 

How to inference?

After learning, inference also resembles the seq2seq model. As shown in the code below, action a is extracted following the addition of the desired return-to-go value to the initial state s. After applying action a to the environment and exiting the environment, the return-to-go in state s' is re-obtained by receiving state s' and reward r, and then placed into DT to extract action a' in the subsequent state s'.

Result

As depicted in Figures 4 and 5, DT, the model proposed in the paper, demonstrated excellent performance on the D4RL and Atari benchmarks, which are widely used in offline RL.

(Figure 4) DT's Performance on the D4RL Benchmark

(Figure 5) DT's Performance on the Atari Benchmark

This study proposes a new paradigm in which the offline RL problem is viewed as a sequence modeling problem. Unlike conventional offline RL algorithms, the model was trained using supervised loss, and a powerful yet complex transformer was successfully implemented.

GENERALIZED DECISION TRANSFORMER FOR OFFLINE HINDSIGHT INFORMATION MATCHING

This research was presented at ICLR 2022 as a generalization and expansion of the current DT. Generalized Decision Transformer (GDT)^[2] proposes a new framework in which DT conditions future trajectory information in policy, and various existing studies can be incorporated into hindsight information matching (HIM) issues.

State marginal matching

The definition of the state marginal distribution, which is the probability of visiting state s when a particular policy π is rolled out, is as follows.

${\mathrm{p}}_t^{\pi }\left(s\right)={\iint }_{s_{0:t},a_{0:t-1}}^{}\prod _t^{}p\left(s_t\overline{)s_{t-1},a_{t-1}}\right)\pi \left(a_t{\overline{)s}}_t\right)$
${\mathrm{p}}^{\pi }\left(s\right)=\left(1-\gamma \right)\sum _{t\text{'}}{\gamma }^{t\text{'}}{\mathrm{p}}_{t\text{'}}^{\pi }\left(s_{t\text{'}}=s\right)$

State marginal matching (SMM) is an objective to find policy π that minimizes the difference between the state marginal distribution ${\mathrm{p}}^{\pi }\left(s\right)$ of the policy and ${\mathrm{p}}^{*}\left(s\right)$ of the target, as defined below.

D is a distribution divergence measure similar to KL divergence, and the SMM objective is primarily used in imitation learning to find a policy to imitate an expert by setting the target as an expert demonstration.

Hindsight Information matching

The author asserts that various studies, such as hindsight experience replay (HER)^[3], temporal difference model (TDM)^[4], learning from play (LfP)^[5], and return-conditioned RL^[1], in which policy that can generate a trajectory that matches future trajectory information are resolving the problem of HIM.

Before defining HIM, it is necessary to define a few additional terms.

- Information statistics $\mathrm{I}\left({\tau }_t\right)$ is a function that represents the statistical property of the partial trajectory ${\mathrm{\tau }}_{\mathrm{t}}=\left(s_t,a_t,s_{t+1},\dots ,s_T,a_T\right)$. For instance, it could be a function that calculates the mean and variance.
- The feature function $\mathrm{\varphi }\left(·,·\right)$ takes the inputs state s and action a and returns a particular feature. For instance, the feature function $\mathrm{\varphi }$ can be the reward function r(s,a) or the function extracted by the sub-dimension of the state.
- When partial trajectory ${\mathrm{\tau }}_{\mathrm{t}}^{\mathrm{\varphi }}=\left(\mathrm{\varphi }\left({\mathrm{s}}_{\mathrm{t}},{\mathrm{a}}_{\mathrm{t}}\right),\dots ,\mathrm{\varphi }\left({\mathrm{s}}_T,{\mathrm{a}}_T\right)\right)$ is provided, which expresses state and action as a feature function, its information statistics are represented as ${\mathrm{I}}^{\mathrm{\varphi }}\left({\mathrm{\tau }}_{\mathrm{t}}\right)$.

 

HIM is a problem that teaches a conditional policy $\mathrm{\pi }\left(\mathrm{a}\overline{)\mathrm{s},\mathrm{z}}\right)$, which receives z as an additional input, to be implemented while satisfying a particular information statistics value z. It is defined as follows.

$\underset{\pi }{\mathrm{m}\mathrm{ax}}{E}_{z~p\left(z\right),r~p_z^{\pi }\left(r\right)}\left[-D\left({\mathrm{I}}^{\mathrm{\varphi }}\left(\mathrm{\tau }\right),z\right)\right]$

Given a trajectory to imitate, by setting s of the trajectory and ${\mathrm{I}}^{\mathrm{\varphi }}\left(\tau \right)$ calculated from the trajectory as z, and training using the input of the conditional policy, the policy can be a policy that follows the trajectory we want to imitate. If feature function $\mathrm{\varphi }$ is used as a reward function and $\mathrm{\gamma }$-discounted summation is used as an aggregate function, information statistics become DT while being defined as ${\mathrm{I}}^{\mathrm{\varphi }}\left({\tau }_t\right)=\sum _{\mathrm{t}}{\mathrm{\gamma }}^{\mathrm{t}}{\mathrm{r}}_{\mathrm{t}}$ at this time. If the feature function $\mathrm{\varphi }$ is a function that outputs the state's position and the aggregate function is a function that merely extracts the end of trajectory ${\mathrm{\tau }}_t^{\varphi }$, then information statistics is ${\mathrm{I}}^{\mathrm{\varphi }}\left({\mathrm{\tau }}_t\right)={\mathrm{\varphi }}_{\mathrm{T}}$, which is and goal-conditioned RL such as are similar to HER. As described previously, based on the feature function and information statistics, various variants can emerge, and this framework is known as Generalized DT. Figure 6 demonstrates how various existing studies can be incorporated into the GDT framework.

(Figure 6) Connection of GDT Framework to Various Previous Studies

(Figure 7) GDT Operation Process Described in this Study

When the target trajectory information statistics are provided to the policy, the author proposes Categorical-DT and Bidirectional-DT, which are new variants of GDT that effectively solve the problems of generating the target trajectory called offline multi-task state-marginal matching (SMM) and offline multi-task imitation learning (IL). Both variants are models that can be obtained by slightly modifying the conventional DT structure. Before introducing the proposed models, we will briefly explain what offline multi-task SMM and IL are.

SMM and IL are distribution matching problems that assess how well the learned policy generates (follows) the target trajectory given the information statistics ${\mathrm{I}}^{\mathrm{\varphi }}\left(\mathrm{\tau }\right)$ of the unseen target trajectory not observed during training. The author defined the SMM problem when the feature function is a specific function, such as a reward function or a function to extract sub-dimensions, and defines the IL problem when the feature function is an identity function ($\mathrm{\varphi }$=s), i.e., when ${\mathrm{I}}^{\mathrm{\varphi }}\left(\mathrm{\tau }\right)$ of the target trajectory is simply given as a full state trajectory.

Categorial DT and Bidirectional-DT

Recently, distributional RL^[6] has demonstrated superior performance to conventional RL in various fields. In general, RL, when learning the Q function, the target, which is the expected discounted reward sum, is used as the scalar value to train the model to make predictions. In contrast, distribution RL is trained to predict the distribution itself as opposed to using the target's expected return value.

CDT is a variant of DT that applies the concept of using the distribution itself rather than scalar values for learning. Information statistics ${\mathrm{I}}^{\mathrm{\varphi }}\left({\mathrm{\tau }}_{\mathrm{t}}\right)$ are defined as $\sum _t{\gamma }^{t}r\left(s_t,a_t\right)$ in DT, a whereas nd as histogram$\left({\mathrm{r}}_{\mathrm{t}},\gamma \right)$ in CDT. Similar to Distributional RL, CDT discretizes a reward with a continuous value, expresses it as a categorical distribution, and trains the model using information statistics.

BDT is a structure that stacked one more layer of transformer on DT. The newly added transformer accepts a reversed state sequence as input and then outputs z, which is used for the transformer policy's information statistics. In other words, while information statistics were computed using predefined feature functions and aggregate functions previously, in BDT, these have been replaced with learnable functions.

Result

The author defines metrics for offline multi-task SMM/IL on the D4RL benchmark, suggests new benchmark tasks, and evaluates CDT and BDT on these benchmark tasks. If the target's state marginal distribution and the learned policy's state marginal distribution are accurately known, it is most accurate to compare the distance via density, but we can only access samples. The author, therefore, uses samples and discretization to approximate the two continuous distributions as a categorical distribution and then calculates the distance.

The following table illustrates CDT's performance in offline multi-task SMM. Figure 8 depicts performance when the feature function is a function that brings x-velocity, a sub-dimension of the state, and Figure 9 depicts performance when the feature function is a reward function. In both cases, it can be confirmed that CDT matches the target distribution better than alternative models.

(Figure 8) Performance of CDT when the Feature Function is x-Velocity

(Figure 9) Performance of CDT when the Feature Function is Reward

To evaluate the generalization performance of CDT, the author also conducts experiments in bi-modal distribution matching and cheetah-velocity matching, which is frequently used in meta RL and IL. To create a trajectory with a bi-modal distribution, the author created a bi-modal trajectory by cutting the running forward trajectory of the cheetah and the backflip trajectory of the cheetah in half and connecting them. Figure 10 depicts the target distribution, the CDT distribution, and the DT distribution. DT only marginally matches the bi-modal distribution, whereas CDT closely matches the target. Figure 11 demonstrates that CDT exhibits similar performance to Meta-BC, which is designed to solve complex multi-task problems.

(Figure 10) Distribution of Sub-Dimension of State in Uni-Modal and Bi-Bodal Trajectory of Target, CDT, and DT

(Figure 11) Comparison Models and CDT Performance in Uni-Modal and Bi-Bodal Trajectories

The cheetah-velocity matching problem is a problem of measuring how well the cheetah robot runs at a desired speed when the trajectories of the cheetah robot running at different speeds are made to run at a speed not observed during learning. The author used trajectories with speeds of 0.5, 1.5, and 2.5 among trajectories with 0.1 intervals between [0.0, 3.0] as the test set, and the remaining trajectories as the train set. As depicted in the table in Figure 12 below, CDT demonstrates superior generalization performance for unknown target trajectories when compared to existing methods.

(Figure 12) Performance of Comparison Models and the CDT in Cheetah-Velocity Matching

Figure 13 below illustrates BDT's performance in offline multi-task IL. In IL, unlike SMM, only the full state of the target trajectory that the learned policy seeks to follow is provided. The author compared BDT and other variants of DT. DT-AE utilizes auto-encoding and DT-CPC is a variant in which feature function $\mathrm{\varphi }\left(\mathrm{s}\right)$ for the state is learned through unsupervised learning using contrastive loss. DT-E2E is a variant in which $\mathrm{\varphi }\left(\mathrm{s}\right)$ is simply trained using the supervised loss of DT.

The author used three strategies when training DT-AE and DT-CPC. The first is a model in which the feature function is simply trained with unsupervised loss, and the second is a model (joint) in which the supervised loss of DT is trained simultaneously. The third model (frozen) is one in which the supervised loss of DT is trained with fixed weights following the unsupervised loss pre-training with unsupervised loss. In addition, all baselines except for BDT aggregate the state sequence output from the feature function is aggregated via summation to produce a single vector, which is then assigned to z of DT. In the case of BDT, an aggregate function, such as summation, is unnecessary because transformer structure is used. BDT matches distributions in IL significantly better than all other variants using learnable $\mathrm{\varphi }\left(\mathrm{s}\right)$.

(Figure 13) BDT Performance in Offline Multi-Task IL

Conclusion

We examined studies that have approached reinforcement learning as a sequence modeling problem using transformer. Research is actively being conducted to create a large-scale generalist model by expanding the aforementioned studies. Existing DT and GDT are limited in that multitasking are only possible in an environment where state and action space are shared. Multi-Game Decision Transformers^[7] successfully trained with a single model in 46 Atari games with varying states and action spaces, and A Generalist Agent^[8] proposed a single model capable of not only Atari games, but also robot control, dialogue, and image captioning. Thus, the boundaries between Vision, NLP, and Reinforcement Learning are blurring, and research into a single large model capable of performing a variety of tasks is underway. Based on the direction of recent research, it appears that human-like AI will be made a reality sooner than expected.