{'fraction_non_alphanumeric': 0.03277781874493031, 'fraction_numerical': 0.010581815500331835, 'mean_word_length': 4.320321694782267, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 0, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Separating the World and Ego Models for Self-Driving

Training self-driving systems to be robust to the long-tail of driving scenarios is a critical problem. Model-based approaches leverage simulation to emulate a wide range of scenarios without putting users at risk in the real world. One promising path to faithful simulation is to train a forward model of the world to predict the future states of both the environment and the ego-vehicle given past states and a sequence of actions. In this paper, we argue that it is beneficial to model the state of the ego-vehicle, which often has simple, predictable and deterministic behavior, separately from the rest of the environment, which is much more complex and highly multimodal. We propose to model the ego-vehicle using a simple and differentiable kinematic model, while training a stochastic convolutional forward model on raster representations of the state to predict the behavior of the rest of the environment. We explore several configurations of such decoupled models, and evaluate their performance both with Model Predictive Control (MPC) and direct policy learning. We test our methods on the task of highway driving and demonstrate lower crash rates and better stability. The code is available at https://github.com/vladisai/pytorch-PPUU/tree/ICLR2022.

In their work, Henaff et al. (2019) develop a model-based approach to policy learning for highway driving. The world model in the proposed system is trained to predict semantic rasterization of the top-down view of a section of a highway around the ego-vehicle, as well as the position and velocity of the ego-vehicle. The policy model interacts with this world model and gets updated by following the gradient calculated by backpropagation from a handcrafted cost through the world model and into the policy parameters. This approach allows for training policies without needing additional on-policy data, which is a very important advantage for self-driving.
However, the proposed world model has a few important limitations. The world model needs to predict both the image of the top-down view and the vehicle state. The two tasks are very different: ego-vehicle state prediction is deterministic and can be computed using a few kinematic equations, while the environment has to be represented by a more complex model capable of representing multimodal predictions to capture the variety of behaviors of other traffic participants. This difference suggests the need for two distinct models, not one. Another limitation of the method proposed by Henaff et al. (2019) is the cost function that is non-differentiable with respect to the position of the vehicle. The gradients flow only through the prediction of the top-down image while omitting important learning signal that can be obtained from the position of the car.
In the present work, we claim that problems that require modeling of agents' motions in a complex stochastic environment should be addressed with two distinct models: a simple and deterministic kinematic model that predicts ego-agents' state in the environment, and a complex model of the environment that is capable of addressing the problem's stochasticity and multimodality. Separating these models provides several advantages. First, we can leverage the knowledge of kinematics to build an exact ego model. Second, when the context allows, we can even make the two models completely independent of each other, making model-predictive control simpler and more efficient.
We build on top of the work of Henaff et al. (2019) and make the following contributions: • we introduce a way to split a world model for self-driving into an environment model and an ego model, and compare different ways of integrating them into one system; • we propose a novel approach that makes the two models independent. We demonstrate that such an approach is faster in certain settings while achieving better performance on the task of highway driving; • we design a cost function that is differentiable with respect to the position and velocity of the ego vehicle;

PROBLEM DESCRIPTION
We consider the problem of autonomous highway driving in this paper, although our approach is applicable to controlling any (similar) autonomous system. We consider three variables at each time step t. They are the self-state s self t ∈ S self , the self-action a self t ∈ A self and the environment state s env t ∈ S env . As the names suggest, the first two are the state of and action taken by the autonomous vehicle under control (ego-car), while the environment state includes everything except the ego-car, such as other cars and any other objects and agents. In our setup where (x t , y t ) are coordinates of the center of the rear axle (approximated as the center of the rear end of the car), (u x t , u y t ) is a unit direction vector, and s t is a scalar denoting the speed. a self t ∈ R 2 is the action taken by the controlled agent at time t, with a t,0 and a t,1 denoting the applied acceleration and rotation strength respectively. s env t is a rasterized mid-level representation of a portion of the highway around the ego-vehicle. An example of such representation is provided in the top row of figure 2. We also assume the availability of a cost function C : S env × S self × A self → R + that, given states and actions at step t, calculates the cost. In this work, C has been handcrafted and includes components to account for proximity of other vehicles, driving off the road and closeness to the lane center. Having all these components, the goal is then to find a policy π that minimizes the cumulative cost Dependencies To create a world model for this problem, we must consider the dependencies among these state and action variables. We start by exhaustively enumerating these dependencies. State components depend on the state information at the previous time step and the action, while the action depends on the state information: The goal of creating a world model then boils down to building a function approximator that predicts s self t+1 and s env t+1 variables given their dependencies. Once we have the models for the states (s self t , s env t ), we can minimize the cost, or the sum of it over time, w.r.t. the action sequence (a self 1 , . . . , a self T ), on which we can train a policy π for driving the autonomous vehicle.

WORLD MODELING
In this section, we present the kinematics model and three possible ways of integrating it with the environment model. In section 4 we present two ways of obtaining a policy using the world model. We then choose 4 combinations of environment model and policy learning set-ups and show experimental results in section 5. The 4 combinations we use are shown in figure 1. We review related work in section 6, and conclude with section 7.
Kinematics model We propose to simplify the dependency pattern of s self t : s self t+1 ← a self t , s self t ,¨s env t . If we assume that the state of the environment s env t does not affect the ego vehicle's state s self t+1 , we can resort to a simple bicycle kinematics model for state prediction. This assumption holds unless there is a collision between the ego car and an object in the environment.
We use the following formulation of the kinematic bicycle model: Where ∆t is the time step (we use ∆t = 0.1 s), unit(v) = v |v| . Equation 4 can be intuitively understood as adding to the current direction vector an orthogonal unit vector multiplied by the turning command and the time step.

Environment model
The environment model f env To find the best way of integrating the kinematic model into the system, we experiment with three configurations of the environment model:

Coupled Forward Model (CFM)
is directly taken from Henaff et al. (2019). In this configuration f env . There is no explicit f self model, instead the world model f env θ is trained to predict both s env t+1 and s self t+1 . The diagram of a set-up using this model is shown in figure 1a. This approach models all dependencies described in section 2 with a single model f env θ .

Coupled Forward Model with Kinematics (CFM-KM)
extends the CFM model, but utilizes the proposed kinematic model and the associated independence assumption to better model s self t . We have f env θ : S env × S self → S env , (s env t , s self t+1 ) → s env t+1 . The model of the environment f env θ , instead of taking a self t as input, takes the prediction of s self t+1 provided by f self (s self t , a self t ). Thus, f env θ does not need to learn the kinematics. Diagrams of two methods using such set-up are shown in figures 1b and 1c.

Decoupled Forward Model with Kinematics (DFM-KM)
In this case, to make the approximation more efficient, we introduce another change to the dependency pattern: This severs the dependency between the environment and the state of the egovehicle. Now, we can run the environment model f env separately from f self . We then have f env . A diagram of a set-up using such pattern is shown in figure 1d.

POLICY
Cost function CFM PL uses the same cost function as Henaff et al. (2019). CFM-KM and CFM-KM MPC add one modification: the off-road component. Adding that cost component to CFM PL cost function does not change performance by much, we show results in appendix C. The resulting cost function C(s env t , s self t ) contains components to account for proximity to other road users, crossing lane demarcations and driving off the road. For the DFM-KM set-up, we implement C km , a modification of the original cost function C that is differentiable with respect to s self . In its original version, the cost is used for backpropagation through s env t and the forward model f env θ (see figure  1a). With the decoupled model, we no longer require the forward model to be differentiable and instead perform backpropagation through s self . To calculate the cost of taking a sequence of actions  to obtain the predictions of (s env 1 , . . . , s env T ) (see the top row of figure 2). Then, we use f self to predict (s self 1 , . . . , s self T ), which are then used to create masks shifted to locations that match the sequence of s self (see the bottom row of figure 2). The mask is shifted in a differentiable manner to allow gradient propagation to s self . The masks are then multiplied with individual channels of the predicted s env to obtain different components of the cost, which are then combined into one scalar with corresponding weighting coefficients. For a more indetail explanation of the cost calculation, see appendix D. With this method, we can backpropagate into the action sequence (a self 1 , . . . , a self T ), update it following the negative direction of the gradients, and repeat the whole process until the cost is low enough. Note that since we do not backpropagate through s env t or f env θ , we do not need to re-run the forward model at each optimization step.
Policy We utilize two approaches for obtaining the driving policy.

Model Predictive Control (MPC)
At step t = 0, having a sequence of planned actions of length T (a self 0 , . . . , a self T −1 ), we want to minimize the cost associated with that plan. We first use forward models f env θ and f self to predict (s env 1 , . . . , s env T ) and (s self 1 , . . . , s self T ), and then use the cost C to obtain the total cost of the predicted trajectory J where γ is the discounting factor that is set to 0.99. Assuming that the cost C calculation is differentiable w.r.t. the actions, we can backpropagate the gradients into the sequence of actions (see figure 1c and 1d). We then do several steps of gradient descent to update the action sequence to minimize the cost. Having the optimized sequence of actions, we then take the first action a self 1 in the sequence and discard the rest, only to re-plan again at the next time step. For more details, see the pseudocode of DFM-KM MPC and CFM-KM MPC in figure 4 in the appendix.

Policy Learning (PL)
This approach is inspired by the method proposed by Henaff et al. (2019). a self t can be modeled using a model of the policy π φ (s env t , s self t ). We train a policy to )·γ t over a roll-out trajectory of T steps. This can be done by computing the negative gradient of the sum of the costs w.r.t. the policy's parameters φ and repeatedly taking the step towards it (see figure 1a).
In the case of both MPC and PL, the proposed sparse dependency pattern simplifies backpropagation. Gradients flow only through the self-states s self t , as we have severed the dependency between the environment state and the self-state. This dramatically lowers the number of interactions involved in the backward pass. This should alleviate the issue of vanishing gradients and improve the quality of the gradients flowing to actions a self t .

Expected benefits
We expect the methods proposed in section 3 to improve the obtained policies in the following ways: 1. We expect the policies to have improved generalization and lower variance of predicted actions when using the kinematics model and/or decoupled forward model. Backpropagating through a simpler and more exact model should enable us to train a more robust policy.
2. Decoupled forward model simplifies the model that connects the action to the cost. Therefore, we expect each backward pass to take less time, making MPC faster than in the coupled approach.
We test both the existence and degree of these benefits in our experiments.
A potential drawback The obvious drawback of the sparsification in the proposed DFM-KM is that the environment state and self-state may eventually become incompatible with each other. Because objects in the environment are not aware of the ego-car, some of them may eventually overlap with the ego-car, resulting in an unrealistic situation, such as squeezing in a traffic jam.
To avoid such unrealistic situations from impacting the policy, we only use a limited roll-out when using the proposed sparse dependency pattern. We argue that this is fine from two perspectives. First, even with the conventional dense dependency pattern, learning a policy by backpropagating through a recurrent network, which is how a world model is often implemented, is challenging because of vanishing or exploding gradients. Therefore, the effective horizon of backpropagation does not decrease much by using a limited roll-out with the sparse dependency pattern. Second, our task of lane following does not require long-range planning by construction. The policy only needs to repeat short-term goals, i.e. to maintain speed and distance from other objects on the road, over and over.

EXPERIMENTS
Dataset We test our methods on the task of highway driving on the NGSIM I-80 dataset (Halkias & Colyar, 2006). The dataset consists of highway driving scenarios recorded from multiple cameras mounted above a section of Highway I-80 in California. The recordings take place at different times of day to maximize the diversity of traffic densities. We follow the pre-processing steps of Henaff et al. (2019) and obtain cars' dimensions and trajectories on the highway. We use the same dataset split of 80%, 10%, 10% for training, validation, and testing respectively.
Crash rates comparison We compare the selected combinations of approaches to policy learning and forward modeling proposed in section 3. The components used by the methods are encoded in the names. For implementation details, refer to appendix B. The compared methods are:   We test all methods in two settings: replay simulation and interactive simulation. Replay simulation simply replays the trajectories of all cars except one, which is controlled by the method we are evaluating. This is the same evaluation protocol as was used by Henaff et al. (2019). This evaluation method has a severe limitation: other cars' actions are simply replayed from the dataset, and therefore are independent of the ego car's actions. Some unrealistic situations may happen, for example, the ego-car can be squeezed by other cars in a traffic jam if the ego car picks a different trajectory from the one that the original vehicle followed during data recording. Since the proposed DFM-KM MPC method assumes exactly such independence, replay evaluation results may be biased in its favor. To address this problem, we also show the results of interactive simulation. Inspired by , we implement interactive simulation by controlling the ego-car with the selected method and controlling all other cars with either the CFM-KM PL or the CFM PL. We do not experiment with controlling other cars with MPC because it is orders of magnitude slower than doing one forward pass with a policy model (see

Variance of actions
We also test if the sparse dependency pattern facilitates more robust predictions by the policies. We show results in table 2. We observe that introducing the kinematic model helps to make the behavior more robust with respect to random initialization, with additional improvement gained from applying the proposed DFM-KM method. For an in-detail explanation of how these variances were computed, see appendix A.

Stress-testing
We test the proposed methods on a hand-designed scenario -controlling an agent on a highway while cruising between two cars, when the car directly in the front brakes suddenly. We show the results in figure 3. We observe that only DFM-KM MPC method manages to successfully complete the scenario. CFM-KM MPC fails, showing that backpropagating through f env θ is not as efficient as backpropagating through f self in DFM-KM MPC approach. We hypothesize that two factors help DFM-KM MPC here. First, in such extreme scenarios the decoupled model has an advantage since it is unreasonable to expect that the car in the front will react to the ego-agent. The independence assumption is justified, and helps the optimization process to find the best action. Second, a learned policy, trained to solve multiple different scenarios, likely becomes a smooth function that cannot take extreme values, while MPC is not restrained by model capacity and can find the minimum of the cost function C better in such extreme cases.

RELATED WORK
Our work is related to the research on model-predictive control and motion planning methods. These are areas with decades of research; for comprehensive overviews of these topics we refer the reader to the books on optimal control (Bryson et al., 1979;Bertsekas, 2005), and motion plan-ning (LaValle, 2006). In this section, we mainly focus on the recent methods that combine the existing approaches with deep neural networks.
Motion planning and Trajectory following Paden et al. (2016) provide a survey of the existing methods for self-driving. Approaches range from graph-search methods, such as A* (Ziegler et al., 2008), to dynamic programming (Montemerlo et al., 2008). In our work, trajectory following is simplified as we are acting in a simulator with a perfectly controllable car, and we focus on motion planning instead.
World modeling has proved to be a great approach to many tasks due to superior sample complexity in policy learning (Nagabandi et al., 2017), although it comes with some caveats, such as the danger of compounding errors (Asadi et al., 2019), and stochasticity (Denton & Fergus, 2018). World model used in our work combines ideas of action-conditioned model (Oh et al., 2015), and stochastic video prediction (Babaeizadeh et al., 2018). Such models have also been used in (Hafner et al., 2020) and (Henaff et al., 2017).
Model Predictive Control has been used widely for self-driving. Zhang et al. (2018) use MPC for planning and collision avoidance, Drews et al. (2018) use MPC to build an impressive system that drives a scaled-down vehicle at high speeds around a track. Our approach can be viewed as a type of Stochastic MPC (Heirung et al., 2018), or Scenario-based MPC (Schildbach et al., 2014;Cesari et al., 2017), where the trajectory is optimized for a limited number of future scenarios (in our case this number is 1, but it can easily be increased).
Kinematic models have been used extensively in self-driving applications. Often, cars are approximated by simplified models, such as unicycle (Kamenev et al., 2021) or bicycle models (Cesari et al., 2017). These models are particularly useful when adding inductive bias to models to produce realistic trajectories in path planning or behavior prediction (Salzmann et al., 2020). Kong et al. (2015) provide an overview of kinematics and dynamics models used for self-driving and apply them to path-following. The work of  is particularly close to ours as they also propose a differentiable kinematic model for training a self-driving policy. However, there is no trained environment model, and the predictions of the other road users are replaced with log replay.
Interactive simulation is a long-standing problem in self-driving cars development.  proposed a method that uses GANs (Goodfellow et al., 2014) to generate the initial state, and then sequentially apply a learned policy to each of the generated agents. Suo et al. (2021) propose a system that models the agents' behavior jointly, making more consistent predictions.

DISCUSSION & CONCLUSION
We presented a novel design of a world model for an agent with known kinematics in a complex stochastic environment. The conducted experiments show that the separation of the world model into an ego model and a model of the environment helps obtain policies that reach better performance in our experiments with highway driving. Decoupling these two models completely makes MPC perform faster and better, and helps to solve a stress-test scenario that requires quick reaction from the policy. We believe that our proposed approach can be applied to any problem that involves an agent with easily predictable kinematics acting in a complex stochastic environment, e.g. controlling robots in the real world, such as delivery carts or drones; or some Atari games, such as Space Invaders, or Freeway. We also believe that for the suggested separation to work, it is not strictly necessary to have the exact kinematic equations, the ego model can also be learned.
There is still more to explore about the proposed approaches. First, although DFM-KM MPC performs better than CFM-KM MPC, it is yet unclear if that is because of the modified cost function C km , or because of the decoupled forward model. Experiments with a method that integrates C km with CFM-KM PL would resolve this ambiguity. Second, comparing the results of CFM-KM PL and CFM-KM MPC, we see that policy learning performs much better, suggesting that DFM-KM PL is an approach worth investigating. Third, hand-designing the cost function is only possible for simple contexts, such as highway driving. For applications to more complex scenarios like urban driving, we would need to learn the cost function, which is highly non-trivial (Ng & Russell, 2000).

B IMPLEMENTATION DETAILS
Forward models To train updated forward models, we follow the procedure proposed by Henaff et al. (2019). The only change to the forward models is the change to the input dimension. The training method was unchanged.

CFM-KM Policy Training
We use the same model as was proposed by Henaff et al. (2019), and train for 70 k steps, with batch size 10, and learning rate of 0.0001. We decrease the learning rate by a factor of 10 after 70% of training.
CFM-KM MPC To find the optimal action, we perform gradient descent for 11 iterations, with learning rate of 0.31. The cost function is calculated as: C = C proximity +0.32·C lane +0.32·C offroad . The uncertainty cost proposed in Henaff et al. (2019) was not used for CFM-KM MPC as it made each iteration impractically slow. We used the plan length of 20 frames, which corresponds to 2 seconds. The hyperparameters were found with random search. We provide pseudo-code in figure 4b.
DFM-KM MPC To find the optimal action, we perform gradient descent for 27 iterations, with learning rate of 0.48. The cost function is calculated as: C = 91.2 · C proximity + 2.88 · C offroad + 3.06 · C lane + 0.1 · C jerk + 0.001 · C destination . We use the plan size of 30. The hyperparameters were found with random search. We provide pseudo-code in figure 4a.
We do not re-train the forward model for this setup, we simply use the CFM-KM, but we always run it with 0-actions -we get predictions s env t+1 that correspond to what the world would have looked like if the ego-vehicle kept going at current speed and did not turn. We then assume that changing the actions would not have changed s env t+1 .

C OFFROAD COST
Another important difference between CFM PL and the other methods is that it does not use offroad cost C offroad . The offroad cost component was introduced to prevent the ego-car from driving off the road. Without it, the cost of driving off the highway is the same as the cost of simply crossing a lane marker. To understand how much the offroad cost contributed to the improvement in performance, we test the CFM PL method with offroad cost. The results are presented in table 3. We see some improvement with adding offroad cost, particularly in interactive evaluation, but the performance does not reach the results of CFM-KM PL and DFM-KM MPC. here is the lane width. These are the distances beyond which the objects are not taken into account in the cost. Then, we want to build a mask that reaches 0 at d x in the front and behind, at d y at the sides, and reaches 1 at the car edges. The masks are: The difference between M side and M car are in the clamping values above 1: when calculating the cost component accounting for proximity to other vehicles, we would like the mask profile to have a "flat nose". α is a hyperparameter used to make the mask non-linear. Higher values make the cost grow more rapidly as objects come closer to the ego-vehicle.
Note that all operations are differentiable with respect to (x, y, u x , u y , s), allowing us to backpropagate through s self .
Calculating cost components Having the masks, we simply perform element-wise multiplication with the corresponding channels of the s env , see figure 5.
C lane = s env lanes , M side (7) C offroad = s env offroad , M side (8) C proximity = s env car , M car (9) (10) The new cost also introduces two new cost components: destination cost, and jerk cost. Destination cost is simply a term that pushes the car to go forward. This prevents cases where MPC fails to drive forward because there are no cars behind. The cost is calculated as C destination = −x. Jerk cost is responsible for making the actions more smooth. Intuitively, this is a cost that penalizes the derivative of the actions w.r.t. time.
. The components are then combined into a single scalar with the corresponding weights α: C =α lane C lane + α offroad C offroad + α proximity C proximity + α destination C destination + α jerk C jerk  Figure 5a depicts the original image, and 5b, 5c, 5d depict channels used for cost calculation. 5e shows the mask used for car proximity cost, while 5f is used for offroad and lane costs respectively.