Sample-efficient Imitative Multi-token Decision Transformer for Generalizable Real World Driving

Abstract

Reinforcement learning via sequence modeling has shown remarkable promise in autonomous systems, harnessing the power of offline datasets to make informed decisions in simulated environments. However, the full potential of such methods in complex dynamic environments remain to be discovered. In autonomous driving domain, learning-based agents face significant challenges when transferring knowledge from simulated to real-world settings and the performance is also significantly impacted by data distribution shift. To address these issue, we propose Sample-efficient Imitative Multi-token Decision Transformer (SimDT). SimDT introduces multi-token prediction, imitative online learning and prioritized experience replay to Decision Transformer. The performance is evaluated through empirical experiments and results exceed popular imitation and reinforcement learning algorithms on Waymax benchmark.

Keywords: Reinforcement Learning, Motion Planning, Autonomous Driving

1 Introduction

The realm of autonomous driving research has witnessed remarkable progress, with simulation technologies [1][2][3][4] reaching unprecedented levels of realism and the burgeoning availability of real-world driving datasets [5][6][7][8]. Despite these advancements, data-driven planning continues to confront a formidable obstacle: the infinite state space and extensive data distribution characteristic of real-world driving.

Imitation learning approaches encounter hurdles [9][10] when presented with scenarios that deviate from the training distribution, exemplified by rare events like emergency braking for unforeseen obstacles. Similarly, these methods grapple with long-tail distribution phenomena, such as navigating through unexpected weather conditions or handling the erratic movements of a jaywalking pedestrian. On the other hand, reinforcement learning (RL) strategies aim to cultivate policies through reward-based learning. RL has difficulty bridging the sim-real gap and sampling efficiency [11]. It often struggle to extrapolate a driving policy that encapsulates the nuanced decision-making process of an experienced human driver, especially when the simulator lacks interactivity or the scenario falls short of realism [12].

Traditional reinforcement learning approaches also struggle with large state space, long-horizon planning and sparse rewards, which are also characteristic of real-world driving scenarios. Decision Transformer [13] leverages a transformer-based architecture to learn policies for decision-making in reinforcement learning tasks via sequence modeling. Despite its potential on scaling with large state space [14], the original architecture and pipeline is deigned for offline learning and is not enough for complex and dynamic autonomous driving task. Classic RL techniques such as prioritized experience replay [15] which is dealing with large scale dataset cannot naturally be applied as Decision Transformer does not compute temporal-difference.

On the other hand, concentrating solely on single-token prediction renders the model excessively susceptible to immediate contextual patterns, thereby neglecting the necessity for more extensive deliberation over protracted sequences. Models trained through next-token prediction methodologies necessitate substantial dataset to achieve a degree of intelligence that humans attain with considerably less token exposure [16]. Receding Horizon Control[17] is a control method that optimizes decision-making over a rolling time horizon, constantly updating its strategy based on newly acquired information. This approach is analogous to multi-token prediction in decision transformers and has potential shifting from myopic to panoramic prediction closed to human cognitive processes.

This paper seeks to address these challenges by proposing an improved Decision Transformer network and a hybrid learning framework that leverages the complementary strengths of imitation and reinforcement learning. Experiment results indicate that our approach yields a substantial enhancement in performance with improvements observed in terms of policy robustness and sample efficiency. The main contributions are as follows:

• We present a fully online imitative Decision Transformer pipeline designed for wide data distribution across large-scale real-world driving dataset.

• We propose multi-token Decision Transformer architecture for receding horizon control to enhance long-horizon prediction and broaden attention field.

• We introduce prioritized experience replay to Decision Transformer and enables sample-efficient training for large-scale sequence modelling based reinforcement learning.

Figure 1: Comparative Illustration of Learning Approaches. The left figure depicts a data distribution of expert data, highlighting its limitations in managing distributional shifts and challenges arising from suboptimal training data. In contrast, the right figure presents our imitative reinforcement learning pipeline that demonstrates enhanced robustness by adapting policies online, thereby achieving superior performance under variable conditions.

2 Related Work

Reinforcement learning via sequence modeling. Trajectory Transformer [18] and Decision Transformer (DT) [13] are pioneer in this area, leveraging transformer architectures to model sequences of state-action-reward trajectories and predicting future actions in offline manner. Following work [19] [20] [21] [22] extends leverage the power of transformers for efficient and generalized decision-making in RL. Online DT [23] and Hyper DT [24] adapt original concept for online settings and interacts with environments. However, previous work are done on relatively simple environments compared to autonomous driving environment.

Multi-token prediction. Transformers have significantly impacted NLP since their inception [25], outperforming RNNs and LSTMs by processing sequences in parallel and efficiently handling long-range dependencies. Subsequent models like GPT [26] and BERT [27] have refined the architecture, enhancing pre-training, fine-tuning, and scalability. Recent studies explore multi-token prediction on semantic representation [28], streamline computation [29], prediction technique [30] and multilingual [31]. However,Focusing only on single-token prediction makes the model too sensitive to immediate context and overlooks the need for deeper analysis of longer sequences[16]. This paper extends the concept to Decision Transformer and explore the potential benefits of multi-token prediction for motion planning.

Learning with real-world driving data. Many work has been done to accommodate with real-world driving data [5][2][6][7] for generalizable driving policy. Lu et al. [32] explores the cooperation between reinforcement learning and imitation learning for real-world driving data. DriveIRL [33] designs an inverse reinforcement learning architecture to learn score component in complex heavy traffic scenarios. TuPlan [9] combines both learning methods with rule-based method for real-world planning. Trajeglish [10] models real world traffic auto-repressively as language processing problem and it has most similar network structure as MSDT. Our approach differs as we apply imitative reinforcement learning for pretraining and online real-world planning adaption for performance enhancement.

3 Methods

In this section, we introduce MSDT, a multi-token sample-efficient reinforcement learning framework via sequence modelling for dynamic driving scenarios. MSDT consists of three components: multi-token decision transformer, online imitative pipeline and prioritized experience replay.

Figure 2: The general network architecture of SimDT. Feature encoding is applied to extract complex real-world driving perceptual data to small but meaningful embeddings. Inside causal transformer, attention relationship are calculated for past context length of $[(s_{t},a_{t},g_{t})]_{T-c}^{T}$. The decoder now predicts multi-tokens for actions and only the first action is applied during the inference stage.

3.1 Network Structure

Since the real-world driving environment is complex and dynamic, specific feature encoding network is designed for the states representation. Real-world driving state contains many perceptual information such as obstacles, road map, traffic and so on. We follow the vectorized representation to organize road map as polylines and then extract with Polyline Encoder [34]. Obstacle with past 10 historical information are recorded in terms of $\bm{[}p_{x},p_{y},v_{x},v_{y},l,w]$. Obstacle and traffic embedding are extracted with multi-layer perception network.

The work further extends the method to goal-conditioned reinforcement learning by adding the relative vector distance between ego vehicle and destination. The importance of goal-condition lies in its influence on the decision-making process of the autonomous agent. Even in an identical environment, the actions taken by the vehicle can vary significantly depending on the specified goal.

Multi-token prediction in causal transformer simultaneously generates multiple tokens in a single forward pass, while still respecting the autoregressive property that ensures each prediction only depends on previously generated tokens. This is typically achieved by using a masked self-attention mechanism that allows the model to consider multiple future positions without violating causal dependencies. As shown in Equation (1), next-token prediction has loss function $L_{a}$ defined as the negative log-likelihood of the policy.

$$L_{a}=-\log\pi_{\theta}(a_{t}\mid s_{t:t-c},a_{t-1:t-c},g_{t:t-c})\$$ (1)

where $\pi_{\theta}$ is the training driving policy. maximize the probability of $a_{t}$ as the next prediction action, given the history of past tokens with context length $c$ of $s_{t:t-c}=s_{t},\,...,\,s_{t-c}.\quad a_{t-1:t-c}=a_{t-1},\,...,\,a_{t-c}.\quad g_{t:t-c}=g_{t},\,...,\,g_{t-c}$.

The loss function is modified for multi-token prediction and assume network predict next 3 tokens. Where $\alpha$ and $\beta$ are the coefficient designed for network to learn more about current step action predictions.

$$L_{ma}= -\log\pi_{\theta}(a_{t}\mid s_{t:t-c},a_{t-1:t-c},g_{t:t-c})$$ (2)

$$-\alpha*\log\pi_{\theta}(a_{t+1}\mid s_{t:t-c},a_{t-1:t-c},g_{t:t-c})-\beta*\log\pi_{\theta}(a_{t+2}\mid s_{t:t-c},a_{t-1:t-c},g_{t:t-c})$$

3.2 Online Imitative Training Pipeline

Algorithm 1 Online Imitative Training Pipeline

Initialize Transition Replay Buffer $D_{trans}$ for capacity A, Trajectory Replay Buffer $D_{traj}$

while $n\leq num\_scenarios$ do

     while $D_{trans}$ is not full do $\triangleright$ Online Data Collection

         if $n\leq 0.5*num\_scenarios$ then

              reproduce scenarios with Human Expert Driving Data, $D_{trans}\leftarrow(s,a,r)$

         else

              reproduce scenarios with Human Expert Driving Data, $D_{trans}\leftarrow(s,a,r)$

              explore scenarios with Policy agent $\pi_{\theta}$ , $D_{trans}\leftarrow(s,a,r)$

         end if

     end while

     HindsightReturnRelabeling: $D_{traj}\leftarrow D_{trans}$, $[[(s_{i,j},a_{i,j},g_{i,j})]_{i=0}^{T}]_{j=0}^{A/T}\leftarrow\left[(s_{i},a_{i},r_{i})\right]_{i=0}^{A}$

     $D_{trans}\leftarrow\emptyset$

     for k in range(1000) do:

         sample and ShuffleObstacleOrder: $[[(s_{i,j},a_{i,j},g_{i,j})]_{t-c}^{t}]_{j=0}^{B}\leftarrow D_{traj}$

         train on sampled data

     end for

end while

The general idea of the proposed algorithm is to perform sample-efficient online imitative reinforcement learning with off-policy expert data for pre-training at beginning. Subsequently, the model undergoes a mixed on-policy adaptation phase which is introduced at the mid-point of the training process. The core concept behind is to quickly shift the distribution towards the expert behavior at beginning and reduce environmental distribution shift with on-policy adaption. Note online adaption and imitative reinforcement learning are performed concurrently after mid of training, this helps the network no to fall into online local minimal.

Imitative reinforcement learning is done by applying similar concept as Shaped IL [35] and GRI [36] where reward is shaped for expert demonstration data. Following same implementation in [32], expert data from real world driving trajectory was converted to expert agent actions with inverse kinematics. We also design negative reward for offroad and overlap (collision) behavior. The network will learn good behavior through imitation reward and bad actions through online interaction with offraod and overlap rewards. The overall online imitative pipeline is essential to achieve the greater data-distributed policy described in Figure 1.

reward function:

$$R_{imitaiton}=\begin{cases}1.0&\text{if }\text{{log\_divergence}}<0.2,\\ 0.0&\text{if }\text{{log\_divergence}}>0.2.\end{cases}$$ (3)

$$R_{offroad}=-2$$ (4)

$$R_{overlap}=-10$$ (5)

However, the real-time collected transition level replay buffer does not contain return-to-go as it can only be calculated after episode is finished and all rewards is collected. Similar to Online Decision Transformer, the transition level replay buffer converted to hindsight trajectory replay buffer when fixed amount of trajectories are collected.

3.3 Prioritized Experience Replay for Decision Transformer

Algorithm 2 Prioritized Experience Replay for Decision Transformer

Initialize Prioritized Trajectory Replay Buffers $D_{single}^{per}$, $D_{overall}^{per}$ with capacity $B$

while $n\leq\text{num\_scenarios}$ do

     Execute lines from Algorithm 1 $\triangleright$ Online Data Collection

     HindsightReturnRelabeling: $D_{traj}\leftarrow D_{trans}$, $[[(s_{i,j},a_{i,j},g_{i,j})]_{i=0}^{T}]_{j=0}^{A/T}\leftarrow\left[(s_{i},a_{i},r_{i})\right]_{i=0}^{A}$

     $D_{trans}\leftarrow\emptyset$

     for k in range(1000) do:

         Sample and ShuffleObstacleOrder: $[[(s_{i,j},a_{i,j},g_{i,j})]_{t-c}^{t}]_{j=0}^{B}\leftarrow D_{traj}$

         train on sampled data and obtain $L_{single}$ and $L_{overall}$

         $D_{single}^{per}\leftarrow\{[[(s_{i,j},a_{i,j},g_{i,j})]_{t-c}^{t}]_{j},L_{single}\}$

         $D_{overall}^{per}\leftarrow\{[[(s_{i,j},a_{i,j},g_{i,j})]_{t-c}^{t}]_{j},L_{overall}\}$

     end for

     Train on $D_{single}^{per}$ and $D_{overall}^{per}$

end while

Prioritized Experience Replay (PER) selectively samples experiences with high temporal-difference errors from the replay buffer for focusing on more informative experiences. However, the Decision Transformer doesn’t use temporal-difference errors, precluding direct application of PER. Instead, we adapt by using action loss to gauge transition importance within the Decision Transformer, which assesses state-action-return relationships. The design concept is that if the model’s predicted actions diverge from actual ones, it indicates a misinterpretation of the environment.

On top of above architecture, extra replay buffers are designed to store prioritised sampled trajectories based on action loss. The action loss represents the difference between the actions predicted by the policy network and the actual actions taken. A low actor loss means that the policy network’s predictions are close to the actual actions, while a high actor loss means that the predictions are far from the actual actions. Prioritised sampled trajectories are stroed based on following criteria:

Criterion 1: Preservation of transitions with maximal single-step action discrepancy: This methodology concentrates on isolating the instances wherein the model’s prognostications manifest the greatest deviation from expected accuracy. Such a strategy is instrumental in directing the model’s learning efforts towards ameliorating its most significant errors.

Criterion 2: Preservation of transitions with maximal cumulative action discrepancy: This methodology is characterized by its emphasis on identifying and retaining sequences wherein the aggregate error of the model’s predictions reaches its apex. This approach holds particular utility for endeavors aimed at refining the model’s performance across a continuum of actions.

The replay buffers store data based on high value in low value out. The prioritized experience replay buffer is sampled for training every fixed amount of episode and its priorities are updated at meantime. The goal for the proposed prioritized experience replay in this paper is to prioritize the trajectories where model has biggest misunderstanding of the corresponding scenarios, and therefore to prioritize on long-tail scenarios.

4 Experimental Results

4.1 Experimental Setup

Agent	Action Space	Sim Agent	Off-Road Rate (%)	Collision Rate (%)	Kinematic Infeasibility (%)	ADE (m)	Route Progress Ratio (%)
Expert	Delta	-	0.32	0.61	4.33	0.00	100.00
Expert	Bicycle	-	0.34	0.62	0.00	0.04	100.00
Expert	Bicycle(D)	-	0.41	0.67	0.00	0.09	100.00
Wayformer[37]	Delta	-	7.89	10.68	5.40	2.38	123.58
BC[38]	Delta	-	4.14$\pm$2.04	5.83$\pm$1.09	0.18$\pm$0.16	6.28$\pm$1.93	79.58$\pm$24.98
BC	Delta (D)	-	4.42$\pm$0.19	5.97$\pm$0.10	66.25$\pm$0.22	2.98$\pm$0.06	98.82$\pm$3.46
BC	Bicycle	-	13.59$\pm$12.71	11.20$\pm$5.34	0.00$\pm$0.00	3.60$\pm$1.11	137.11$\pm$33.78
BC	Bicycle(D)	-	1.11$\pm$0.20	4.59$\pm$0.06	0.00$\pm$0.00	2.26$\pm$0.02	129.84$\pm$0.98
DQN[39]	Bicycle(D)	IDM	3.74$\pm$0.90	6.50$\pm$0.31	0.00$\pm$0.00	9.83$\pm$0.48	177.91$\pm$5.67
DQN	Bicycle(D)	Playback	4.31$\pm$1.09	4.91$\pm$0.70	0.00$\pm$0.00	10.74$\pm$0.53	215.26$\pm$38.20
SimDT(ours)	Bicycle	Playback	3.36$\pm$0.04	2.65$\pm$0.06	0.00$\pm$0.00	6.73$\pm$0.41	-

Table 1: Performance evaluation are done against IDM simulation agents. Agents run without any termination conditions in WOD1.2 evaluation dataset. Action space is continuous unless denoted with D (discrete). Waymax benchmark table is used as we use exact same experiment settings.

Dataset, simulator and metrics. Training and Experiments are done based on Waymo Open Dataset and Waymax simulator. Waymax provides embedded support for reinforcement learning and diverse scenarios drawn from real driving data. Waymax incorporate with Waymo Open Motion Dataset (WOMD) which provides $531,101$ real-world driving scenarios for training and $44,096$ scenarios for validation, each scenario contains 90 frames of data. Specifically, WOMD v1.2 and exact same metrics (off-road rate, collision rate, kinematic infeasibility, average displacement error (ADE)) from Waymax are used to benchmark with the paper.

Implementation Detail. Models of various sizes are developed to quickly conduct ablation studies and assess final performance effectively. Raw observation takes nearest ego vehicle, 15 nearest dynamic obstacles, 250 of closest roadgraph elements, traffic signals and position goal as input. The total size for each step observation is $7050$ and feature extraction is applied to reduce the total size. SimDT(tiny) has 256 tokens for each element of $(s,a,g)$ pair, 6 blocks, 16 attention head and in total 7.7 million parameters. SimDT(small) has 384 tokens for each element of $(s,a,g)$ pair, 10 blocks, 16 attention head and in total 22.2 million parameters. Both models use context length with value 10, which means causal transformer has access to past 10 $(s,a,g)$ pairs.

4.2 Benchmark Comparison

SimDT is evaluated using Intelligent Driving Model (IDM)[40] as the simulated agent. SimDT achieves Off-Road Rate of 3.36%, Collision Rate of 2.65%, Kinematic Infeasibility of 0.00%, and ADE of 6.73m. SimDT significantly outperforms them in collision rate and being second in off-road rate against other learning-based approaches. Compared same reinforcement learning category method, SimDT demonstrates a substantial reduction in Off-Road Rate and Collision Rate than DQN. Suggesting that our method is more effective at keeping the vehicle on the road and avoiding accidents. The Off-Road Rate of SimDT is higher than the best performing BC ’Bicycle (D)’ model by 2%. Similarly, the Collision Rate of SimDT shows a 1.94 percentage point improvement over the same BC model. This improvement in safety-critical metrics highlights the robustness of SimDT in real-world driving scenarios.

When compared to expert demonstrations, SimDT achieves competitive results in terms of safety metrics Collision Rate are within the same magnitude as those reported by the experts. However, the ADE of SimDT is notably higher at 6.73m, which is approximately 6 meters away from the expert models. This suggests that SimDT learns a safe and feasible policy but different from the expert recording. While the ADE for SimDT is higher than that of other imitation learning models, it is important to note that ADE alone may not capture the complete picture of driving performance. The emphasis on safety and kinematic feasibility by SimDT may contribute to a cautious driving style, which can result in a slightly higher ADE but with significantly safer outcomes.

4.3 Ablation Study

Agent	Off-Road Rate (%)	Collision Rate (%)	Kinematic Infeasibility (%)	ADE (m)
DT(tiny)	5.86	3.43	0.00$\pm$0.00	7.64
DT(tiny) + PER	4.28	3.27	0.00$\pm$0.00	7.26
DT(tiny) + PER + OPA	3.73	2.79	0.00$\pm$0.00	7.05
DT(tiny) + PER + OPA + 2 token prediction	3.74	2.72	0.00$\pm$0.00	7.09
DT(tiny) + PER + OPA + 3 token prediction	3.57	2.54	0.00$\pm$0.00	6.97
DT(small) + PER + OPA + 3 token prediction	3.36	2.65	0.00$\pm$0.00	6.73

Table 2: Ablation Study. PER is prioritized experience replay, OPA is online-policy adaption.

(a) Irregular Traffic Junction

(b) Parking Slot

(c) T Junction

Figure 3: Illustration of Data Selection for Prioritized Experience Replay: 3(a) is chosen because its uncommon expert behavior that need to slow down while steer to the right to keep lane. 3(b) illustrates a rare parking situation and highlights a case that was picked because it had the most mistakes when looking at the whole series of actions. 3(c) is kept as the suboptimal action taken in that situation was not reproduced given the corresponding low return-to-go.

Prioritized Experience Replay for Decision Transformer. Since our proposed imitative reinforcement learning can obtain almost infinite amount of dataset through online interaction, the ability of prioritized experience replay becomes critical for sample efficiency. Compare to pure Decision Transformer, the model which adapts PER has 1.58% and 0.16% reduction in off-road and collision rate. Decision Transformer with PER is able to reach same performance with 80% of data. There are three types of the data that is preferentially stored for PER (Fig. 3). The initial category encompasses instances wherein a discernible discrepancy arises between the predicted actions of the learning model and those executed by an expert. The second category pertains to scenarios wherein the cumulative action loss associated with a particular trajectory is substantially elevated, a phenomenon that predominantly transpires within the confines of rare encountered environmental conditions. The third category is representative of situations where trajectories indicative of suboptimal online adaptation are documented, highlighting the model’s challenges in identifying and rectifying suboptimal behaviors. The sample-efficient leanrnig curve can be found in Appendix. 6

Multi-token Decision Transformer. Due to physics limitation of real world vehicles such as inertia and momentum, actions taken at current time-step can significantly affect the following time-step actions. Current state has effect to near future actions steps. eg. reckless pedestrian crossing can cause emergency breaking for ego vehicle and it takes at least few steps to finish. It is important for Neural Network to understand return-to-go and sequence action consequence during the training stage. On the other hand, Focusing only on predicting one token at a time makes a model too sensitive to the immediate context, overlooking the need to consider longer sequences of token for better understanding. Multi-token prediction allows for a more nuanced grasp of world interaction with less data. Compared with single-token prediction, 3-token SimDT has 0.16% and 0.25% improvement in off-road rate and collision rate.

(a) scenario

(b) 1-token prediction

(c) 3-token prediction

Figure 4: Attention map comparison for single-token and multi-token prediction. Multi-token prediction network has more diverse attention field.

(a) demonstration(navy) and exploration trajectory

(b) learned policy

Figure 5: Illustration of how demonstration and exploration trajectory to learn a generalized policy.

5 Conclusion and Discussion

We introduces SmiDT, an innovative approach to sequence modeling based reinforcement learning, particularly targeted for the complexities of real-world driving scenarios. Our fully online imitative Decision Transformer pipeline is adept at handling diverse data distributions found within extensive driving datasets, ensuring wide applicability and robustness. By implementing a multi-token Decision Transformer that integrates receding horizon control, we improve the model’s ability to predict over longer horizons and extend its attention span across broader contexts. Furthermore, the incorporation of prioritized experience replay within our framework enhances the sample efficiency of training, allowing for more effective learning from large-scale datasets. Our work can also benefit other real-world robotics tasks that demand sample-efficient imitative reinforcement learning.

Limitation. Due to computational constraints, we couldn’t train a larger network with increased embedding sizes, more transformer blocks, additional attention heads, and extended context length. A longer context would enhance the model’s grasp of its environment, potentially improving its capability for both high-level task planning and low-level action planning.

Acknowledgments

If a paper is accepted, the final camera-ready version will (and probably should) include acknowledgments. All acknowledgments go at the end of the paper, including thanks to reviewers who gave useful comments, to colleagues who contributed to the ideas, and to funding agencies and corporate sponsors that provided financial support.

Appendix

A. Learning cure

where model 1 is Pure Decision transformer, and model 2 has Decision Transformer with Prioritized Experience Replay. Figure shows our Sample-efficient Imitative Pipeline converges faster and has better performance.

Figure 6: Learning Curve. where Model 1 is Pure Decision transformer, and Model 2 has Decision Transformer with Prioritized Experience Replay. Figure shows our Sample-efficient Imitative Pipeline converges faster and has better performance.

B. Metrics Definition

Collision rate This metric checks for overlap between bounding boxes of objects in a 2D top-down view at the same time step to determine if a collision has occurred.

Off-Road rate indicates the percentage whether the vehicle is driving within the road boundaries, with any deviation to the right of the road’s edge considered off-road.

Kinematic Infeasibility Metric is binary metric assesses whether a vehicle’s transition between two consecutive states is within predefined acceleration and steering curvature limits, based on inverse kinematics.

Average Displacement Error (ADE) calculates the mean L2 distance between the vehicle’s simulated position and its logged position at corresponding time steps across the entire trajectory.

Route Progress Ratio calculates the proportion of the planned route completed by the vehicle, based on the closest point along the path at a given time step. Route Progress Ratio feature is not released yet and benchmark in this paper will skip this metric.