Sequential Action-Induced Invariant Representation for Reinforcement Learning

SAC is an off-policy reinforcement learning algorithm for continuous control, which consists of a Critic network for learning the value function $Q_{\varphi}(o;a)$ with parameters ${\varphi}$ and an Actor network for learning the policy function $\pi_{\psi}(a|o)$ with parameters ${\psi}$. Unlike value-based methods, SAC optimizes a stochastic policy to maximize the expected trajectory reward. The main highlight of SAC is that the maximum entropy item ${\cal H}=log(\pi_{\psi}(a|o))$ is added to decentralize sampling actions, which enhances the exploration and robustness of the algorithm.

Specifically, The Critic network first samples transition $e_{t}=(o_{t},a_{t},o_{t+1},r_{t},d)$ from the replay buffer $\cal D$ to minimize the Bellman error, where $d$ represents done signal. The training of parameters $\varphi$ can be represented by the following loss:

The target $\cal T$ computes the expectation of the next actions sampling from current policy, defined as:

where network $\hat{\varphi}$ comes from the Exponential Moving Average (EMA) of the Critic parameter $\varphi$, and the parameter $\alpha$ is a positive entropy coefficient that determines the priority of entropy maximization over value function optimization.

Finally, we sample actions $a\sim\pi_{\psi}$ from the policy with parameter $\psi$, and train the Actor by maximizing the expected reward of the sampled actions:

Behavior Similarity Metrics (BSM) related methods generally employ MDP tuples, such as reward, action, policy, etc., to measure the task equivalence between observations, where an auxiliary loss is usually constructed to promote the extraction of task-relevant state representations. Policy Similarity Metric (PSM) is a BSM-based method that utilizes policy distance. Since our work is closely related to PSM in terms of optimization, we briefly introduce the basic definition of PSM and the differences between them.

For the same point, both our method and the PSM method are based on the BSM principle, i.e., if the encoder distance is equal to the distance between certain MDP tuples (such as the PSM distance in Theorem 3.1), it is considered that all the information extracted by the encoder is task-related. For differences, PSM method optimizes the one-step distance $\lVert z_{i}-z_{j}\rVert_{1}$ of the encoded states, making it close to the one-step distance between the MDP tuples that can determine the task, i.e., the sum distance of the policy distance $DIST(\pi^{*}(x),\pi^{*}(y))$ with the next transition distance ${\cal W}_{1}$. However, in our method, besides the basic idea of sequence actions, the most important difference is that we minimize the difficult multi-step distance, while there is no need to learn the transition model ${\cal P}$ avoiding the impact of model accuracy.

The construction method of auxiliary tasks tremendously determines the expected performance. Therefore, during the construction of auxiliary tasks, a key step is how to construct auxiliary objectives that can efficiently optimize the state encoder. To that end, we construct an auxiliary task with optimal action sequences and make them guide the extraction of task-relevant information. In detail, in order to take full advantage of the properties possessed by the action sequence, we establish a probability model for sequence actions and finally employ the characteristic function of the probability distribution as the target of the auxiliary loss. Before this implementation, we first use the newest policy to calculate the actions of sampling observations and try to find better action sequences.

Actions Sequence Sequential action-induced representation aims to lock task-relevant information in observations that have a specific mapping relationship with action sequences, so action sequences will directly guide the extraction of task-relevant information. In theory (Theorem 3.1), SAR should also utilize the optimal action sequences sampled from optimal policy $\pi^{*}$, which helps to learn representations accurately Agarwal et al. (2020). As for reason, we empirically believe that due to the optimal sequential actions sampled from an optimal probability distribution function (PDF), the information controlled by optimal actions is deeply decoupled from the random noise information. In contrast, information controlled by random actions cannot be accurately extracted due to its coupling with random noises.

In practice, the optimal actions cannot be obtained during training, but the actions sampled from the suboptimal policy can also be worked in representation learning with an accuracy limitation Agarwal et al. (2020). Therefore, to further improve the optimality of action sequences, we make the following improvement: At each training step, for the action $a_{t}$ sampled from the replay buffer, we do not directly use the $a_{t}$ that was calculated with old policy $\pi^{old}$ like Agarwal et al. Instead, we update the $a_{t}$ by the latest current policy. To be specific, we first sample sequential observations $\{o_{t},o_{t+1},...,o_{t+T}\}$ of length $T$ from replay buffer, and then use the current policy $\pi^{new}$ to calculate the action sequence with the sampled sequential observations:

Compared with the old policy adopted by Agarwal et al., current policy $\pi^{new}$ is thought to be closer to the optimal distribution of $\pi^{*}$ (see Assumption 4.1), so adopting the $\pi^{new}$ will guarantee a tighter lower bound of the action optimality. For easier reading, suboptimal policy $\pi^{new}$ and suboptimal action $a^{new}$ ($a^{new}\sim\pi^{new}(\cdot)$) will be marked with $\pi^{*}$ and $a^{*}$ respectively below.

Next, the characteristic function of action sequence distribution is used as the target $\cal{Y}$ of the auxiliary loss in the modeling process and calculation formula.

Consider a set of sequential actions $a^{s}=(a^{*}_{t},a^{*}_{t+1},...,a^{*}_{t+T})\in\Psi$ under $\pi^{*}$, where $\Psi$ is the combined space of the T-step actions defined on original action space $\cal A$, and the probability distribution function of $a^{s}$ is expressed as $f_{\Psi}(a^{s})$. Then, the characteristic function defined under $\Psi$ is expressed as follows:

Where imaginary unit $i=\sqrt{-1}$, inner product $<\theta,a^{s}>={\sum}^{k=t+T}_{k=t}\theta_{k}a^{*}_{k}$, and $\cal Y$ is the complex function of the real variable $\theta\sim\Theta$. $\Theta$ is the probability density function on $\mathbb{R}^{T}$, where Gaussian distribution $\Theta={\cal N}(\mu,\sigma^{2})$ is actually adopted.

Finally, the characteristic function of action sequence distribution can be expressed by the real and imaginary units as follows:

In this subsection, we mainly introduce the implementation and optimization details of the auxiliary task and give the overall training algorithm of the SAR framework.

First, we estimate the characteristic function of the action sequence distribution using the neural network predictor $Pred$. To improve the generalization of representation learning, the input of the $Pred$ consists of a concatenation of encoding state $z$, real variable $\theta$, and reward sequence $r^{s}=\{r_{t},r_{t+1},..,r_{t+T}\}$. Therefore, the output of $Pred$ prediction can be denoted as:

Second, we compute the $L_{2}$ distance (mean squared error) between the predicted characteristic function and the true characteristic function of the action sequence, then optimizing the encoder $\phi$ through preliminary loss function:

In fact, since the true characteristic function cannot be directly obtained, following Yang et al. Yang et al. (2022), we optimize an upper bound of ${\cal L}^{MSE}(\cal G,Y)$:

In Eq.12, since the true characteristic function is represented by real and imaginary units, the predicted vector is also decoupled into two corresponding parts $[{\cal G}_{\phi}^{cos},{\cal G}_{\phi}^{sin}]=\cal G_{\phi}$. Finally, the loss function with decoupled real and imaginary parts is expressed as:

Additionally, we construct a cosine similarity (CS) loss function on the cosine component between the true characteristic and predicted vector. Unlike $L_{2}$ distance, which optimizes the magnitude of the prediction vector, the cosine similarity loss can guarantee the phase distance between the predicted and the target vector, thus making the prediction vector closest to the characteristic vector of the true action sequence.

In the end, the above auxiliary representation learning module with ${\cal L}^{MSE}$ and ${\cal L}^{CS}$ will be extended to the Soft Actor-Critic (SAC) reinforcement learning framework. And the AC framework includes the Actor loss ${\cal L}^{\pi}$ for the training policy and the Critic loss ${\cal L}^{V}$ for the training value network. Therefore, we train SAR through minimizing the total loss ${\cal L}={\cal L}^{MSE}+{\cal L}^{CS}+{\cal L}^{\pi}+{\cal L}^{V}$, and corresponding pseudocode can be seen in Algorithm 1.

DMControl with Background Distraction The DMControl suite Tassa et al. (2018) is a challenging visual environment based on the Mujoco physics simulator, which is widely used for the performance verification of visual reinforcement learning. Following to the distracting DMControl setting Stone et al. (2021), the background distraction of the experiment is sampled from natural videos in the DAVIS 2017 dataset. Specifically, algorithms use 2 videos alternately as background during the training phase and evaluation on unseen 30 videos. We adopt the above challenging settings to evaluate the anti-distraction representation ability of algorithms.

Baselines We compare SAR against recent strong visual RL baselines under the DMC suite: (i) CRESP Yang et al. (2022), an algorithm that proposes to induce task-relevant state representations through reward sequences, is one of the current algorithms with outstanding generalization performances; (ii) DrQ Yarats et al. (2021), a recent state-of-the-art algorithm that utilizes data augmentation techniques to learn robust representations from pixels; (iii) DBC Zhang et al. (2020) which is a common RL baseline for learning compact encodings based on the task equivalence principle under bisimulation metric; (iiii) CURL Laskin et al. (2020a), a contrastive unsupervised representation learning method for RL, which has achieved state-of-the-art data-efficiency on clean pixel-based environments.

The evaluation results of the experiment on the DMControl tasks with the unseen distracting background are shown in Fig 3. In general, we evaluate the performance of SAR by comparing with four common baselines: DrQ, DBC, CURL, as well as the latest CRESP, in 6 tasks of different difficulty of Cartpole, Cheetah, Hopper, and Ball_in_up in four scenarios. In these experiments, all baselines adopt the settings consistent with the original work. For SAR, the learning rate is 5e-4, the initial temperature is 0.1, the update frequency of Critic and Actor is 2, and we use Adam optimizer Kingma & Ba (2014) to optimize the auxiliary task, Critic and Actor network with a batch size of 256. See Table 3 for the full of hyperparameters.

In order to highlight that our method has a stronger representation ability for eliminating background distractions, we choose 4 difficult tasks (Ball-in-cup-catch, Cartpole-swingup, Cheetah-run, and Reacher-easy) from the DMC500K benchmark, and additional 2 high-difficulty tasks with sparse reward property (Hopper-stand, Cartpole-swingup-sparse). As shown in Table 1 based on 5 repeated experiments, the SAR method with sequential action-induced task information representation achieves the best performances on these tasks compared to the baselines.

We quantify and analyze the ability of different representations to induce task-relevant information during the training and evaluation stages. Fig. 4 shows the performance comparison curves of baselines during the training stage. In this stage, with iterative training, the distracting backgrounds (videos) will be gradually learned and familiarized by agents, i.e., the distracting backgrounds are seen for agents, so the agents are easier to achieve generalization. In contrast, since the distracting backgrounds used during training and evaluation are independent of each other, i.e., the background videos used during the evaluation stage are not seen during the training stage, some models with weak generalization ability may have significant performance degradation due to the unseen distracting backgrounds. Of course, if the representation learner can sufficiently understand and determine the task-relevant information, the agent should perform equally in the training and evaluation stage.

To quantify the generalization ability of the models, we denote the generalization decay ratio of the model performance from the training (seen background distractions) to evaluation (unseen background distractions) stage as $\rho=\frac{{score}_{train}-{score}_{evel}}{{score}_{train}}$. The comparison bar charts of the generalization decay ratio for each method on different tasks are shown in Fig. 5. To be specific, we can summarize the main two conclusions: i) Our SAR method that learns state representation via action sequences achieves the lowest generalization decay ratio compared with the baselines on most tasks, and shows better robustness to the unseen video backgrounds. ii) There is a performance loss when the methods generalize from training to the evaluation stage, and the general decay range of performance is $10\%\sim 30\%$, where the largest performance loss is on the complex tasks of Hopper-stand and Cartpole-swingup-sprse.

We employ the t-SNE visualization to observe the distribution maps of observation transitions under the latent state spaces of SAR and DrQ encoders, where the transitions come from the same batch of 4-episode interaction data, and to establish the relationship between the distribution and the observations with background distractions. In addition, we associate observation images and embedding points according to their sequence numbers annotated in advance. As shown in Fig. 6, the embedding maps on the left and right correspond to the 2D projections of the observations in the SAR and DrQ latent state spaces, respectively, and the middle shows three groups of random observation samples with similar task information but different video backgrounds. For the projection pairs of these samples, we find that they are very close in the latent state space of our SAR and have similar state values (similar colors). Contrary to SAR, the distances and state values in the latent spaces of the DrQ encoder are far apart, which obviously deviates from reality. In other words, for unseen background distractions, the proposed SAR method can accurately extract task regions through action sequences and compresses them into a low-dimensional representation, which has task and value information that are nearly consistent with the real state. In summary, the t-SNE visualization experiment strongly confirms our previous hypothesis that the method of sequential action-induced invariant representations can better learn representations than baselines.

The inputs of real-world visual control systems such as autonomous driving often contain task-irrelevant elements, such as clouds, mountains, and light, it is necessary to extract as much task-relevant information as possible, e.g., the information about road parts, other vehicles, or obstacles, in order to better autonomous driving under complex and unstructured distractions. To verify the state representation and generalization ability of the algorithms under real visual observations, we applied them to the field of visual autonomous driving based on the CARLA simulator Dosovitskiy et al. (2017).

CARLA Task CARLA is a powerful autonomous driving simulator that has visual inputs and physical effects nearly the same as the cars in the real world, where we can set rich task components, weather, lighting and NPC vehicles in the scene to meet requirements of different tasks. In this experiment, we choose the officially provided Town04 map with a ring highway containing a crossroad, as shown in Fig. 7. Concretely, we set up a 300-degree wide-angle image collector composed of 5 cameras with 60-degree views on the roof of the vehicle agent. The agent needs to operate the steering, brake and accelerator to control the vehicle’s movement under conditions of wind, rain, cloud or sunlight changes.

Task Setting  The goal of agents is to safely drive as far as possible in a limited 1000 environmental steps, and the episode will be reset when the agent exceeds the maximum environmental step, goes out of road bounds, or crashes. The observation input for training is an RGB image concatenated by 5 cameras, with a size of $84\times 420$ pixels, and the control actions are standardized speed and steer. Following Zhang et al. Zhang et al. (2020), the reward function is set as $r_{t}(s,a)=\textbf{v}^{T}\hat{\textbf{u}}_{highway}*\Delta t-C_{i}*impulse-C_{s}*|steer|$, where $\hat{\textbf{u}}_{highway}$ is the unit vector of the highway, $\textbf{v}^{T}\hat{\textbf{u}}_{highway}$ represents the effective speed which is projected by the vehicle speed onto the highway, and then is multiplied by the unit time $\Delta t$ to represent the effective distance, $impulse$ is the clash force ($N/s$) obtained via physical calculation, and $steel$ is the output amplitude of the steering. Detail parameters can be seen in Table 4. Finally, it should be emphasized that learning a basic driving policy through very short training steps (1e5 steps) in such real complex scenarios is a challenging task.

Evaluation  We choose CRESP, CURL, and DrQ algorithms as the baselines. As shown in Fig. 8, our algorithm is significantly better than the baselines in both crucial comparisons of episode reward and episode distance. To be specific, the best average reward and driving distance of the proposed algorithm are 193 $m$ and 219 $m$, which are 29.5% and 27.3% higher than the suboptimal baseline respectively, as shown in Table 2. In addition, although the learning parameters of our algorithm are slightly increased than those of the CURL and DrQ, the convergence rate is still higher than that of all baselines. The successes episode ratio is another important metric, which represents the proportion of episodes that reached the episode horizon. As can be seen from Table 2, our method also improves by 42.9% over the suboptimal baseline. Furthermore, for the metrics of the average steering amplitude change, average brake degree, and crash intensity, which represent the stability metrics of the vehicle during driving, our method outperforms the baselines in terms of the first two metrics, for crash intensity, the performance of our algorithm is lower than the baselines, the possible explanation is as the agents of the baseline travel a shorter distance than ours, so some of the crashes can be avoided. As a whole, the significant advantages of the comparative experiments empirically indicate that the proposed algorithm is able to better handle task-irrelevant distractions, thereby learning an excellent policy.