SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Alberto Silvio Chiappa
École Polytechnique Fédérale de Lausanne (EPFL)
alberto.chiappa@epfl.ch
Alessandro Marin Vargas
alessandro.marinvargas@epfl.ch
Ann Zixiang Huang
Mila, EPFL
zixiang.huang@mail.mcgill.ca
alexander.mathis@epfl.ch
In Reinforcement Learning, agents learn policies by exploring and interacting with the environment. Due to the curse of dimensionality, learning policies that map high-dimensional sensory input to motor output is particularly challenging. During training, state of the art methods (SAC, PPO, etc.) explore the environment by perturbing the actuation with independent Gaussian noise. While this unstructured exploration has proven successful in numerous tasks, it can be suboptimal for overactuated systems. When multiple actuators, such as motors or muscles, drive behavior, uncorrelated perturbations risk diminishing each other's effect, or modifying the behavior in a task-irrelevant way. While solutions to introduce time correlation across action perturbations exist, introducing correlation across actuators has been largely ignored. Here, we propose LATent TIme-Correlated Exploration (Lattice), a method to inject temporally-correlated noise into the latent state of the policy network, which can be seamlessly integrated with on- and off-policy algorithms. We demonstrate that the noisy actions generated by perturbing the network's activations can be modeled as a multivariate Gaussian distribution with a full covariance matrix. In the PyBullet locomotion tasks, Lattice-SAC achieves state of the art results, and reaches  $18\%$  higher reward than unstructured exploration in the Humanoid environment. In the musculoskeletal control environments of MyoSuite, Lattice-PPO achieves higher reward in most reaching and object manipulation tasks, while also finding more energy-efficient policies with reductions of  $20 - 60\%$ . Overall, we demonstrate the effectiveness of structured action noise in time and actuator space for complex motor control tasks. The code is available at: https://github.com/amathislab/lattice.
Effectively exploring the environment while learning the policy is key to the success of Reinforcement Learning (RL) algorithms. Typically, exploration is attained by using a non-deterministic policy to collect experience, so that the random component of the action selection process allows the agent to visit states that would have not been reached with a deterministic policy. In on-policy algorithms, such as A3C [1] and PPO [2], the policy network parametrizes a probability distribution, usually an independent multivariate Gaussian, from which actions are sampled. In off-policy algorithms, such as DQN [3], DDPG [4] and SAC [5], the policy used to collect experience can be different from the one learned to maximize the cumulative reward, which leaves more freedom in the selection of an exploration strategy. The standard approach in off-policy RL consists in perturbing the policy by introducing randomness (e.g.,  $\epsilon$ -greedy exploration in environments with a discrete action space and Gaussian noise in environments with a continuous action space). In the past years, many notable
A LATTICE - LATent TIme-Correlated Exploration
Tne aee
Figure 1: A Lattice introduces state-dependent perturbations both to the action space and to the latent space. B Compared to the default action noise, Generalized State-Dependent Exploration (gSDE) [8] introduces a state-dependent perturbation of the action. Lattice generalizes gSDE, including a second, policy-dependent perturbation, which induces correlation between action components. C Similarly to gSDE, the perturbation matrices  $P_{\mathbf{x}}$  and  $P_{\mathbf{a}}$  are sampled periodically, providing noise with a temporal structure.
successes of RL are based on independent noise [1, 2, 4, 6, 7]. However, previous work has challenged the inefficiency of unstructured exploration by focusing on introducing temporal correlations [8, 4, 9-11]. Indeed, two opposite perturbations of subsequent actions might cancel out the deviation from the maximum probability trajectory defined by the policy, effectively hiding potentially better actions if followed by a more coherent policy. Here, we argue that not only correlation in time [4, 9-11, 8], but also correlation across actuators, can improve noise-driven exploration.
In motor control, precise coordination between actuators is crucial for the execution of complex behaviors [12-15]. Perturbing each actuator independently can disrupt such coordination, limiting the probability of discovering improvements to the current policy. In particular, musculoskeletal systems [16-19] typically feature a larger number of actuators (muscles) than degrees of freedom (joints). Here, we demonstrate that in such systems the exploration achieved by uncorrelated actuator noise is suboptimal. To tackle this problem, we introduce an exploration strategy, named LATent TIme-Correlated Exploration (Lattice). Lattice takes advantage of the synergies between actuators learned by the policy network to perturb the different action components in a structured way. This is achieved by applying independent noise to the latent state of the policy network (Fig. 1 A). With extensive experiments, we show that Lattice can replace standard unstructured exploration [2, 5] and time-only-correlated exploration (gSDE) [8] in off-policy (SAC) and on-policy (PPO) RL algorithms, and improve performance in complex motor control tasks. Importantly, we demonstrate that Lattice-SAC is competitive in standard benchmarks for continuous control, such as the locomotion environments of PyBullet [20]. In the Humanoid locomotion task, Lattice improves training efficiency, final performance and energy consumption. In the complex object manipulation tasks of the MyoSuite library [18], involving the control of a realistic musculoskeletal model of the human hand, Lattice-PPO achieves higher reward, while also finding more energy-efficient policies than standard PPO. To the best of our knowledge, our work is the first to showcase the potential of modeling the action distribution of a policy network with a multivariate Gaussian distribution with full covariance matrix induced by the policy network weights, rather than a diagonal covariance matrix.
B
Figure 2: A Representation of a musculoskeletal model of a human arm. The elbow joint is actuated by flexor-extensor muscle groups. Figure adapted from [18, 31]. B Covariance matrix of the actions (left) and noise of the actions (right) when latent state of a learned policy is perturbed for flexors  $f_{1} - f_{3}$  and extensors  $e_{1} - e_{3}$ . C Distribution of the joint angle perturbation across episodes when noise is applied to the action or to the latent space. The action noise was tuned so that each action component has the same variance as with latent noise. The difference in distribution, due to the extra-diagonal terms of the action covariance, is statistically significant ( $p = 3.2 \times 10^{-17}$ , Wilcoxon signed-rank test).
Exploration based on the learned policy. In on-policy RL, the target policy must be stochastic, so that the experience collected includes enough variability for the agent to discover policy improvements. Algorithms such as A3C [1], PPO [2], and TRPO [21] parametrize each action component with an independent Gaussian distribution (continuous actions) or a categorical distribution (discrete actions). In off-policy RL, most exploration strategies are based on perturbations of the policy. For example, in discrete action spaces an exploratory policy can be derived from the target policy by adding  $\epsilon$ -greedy exploration, or more sophisticated perturbations such as the upper confidence bound [22, 23], Thompson sampling [24], Boltzmann exploration [25] and Information-Directed Sampling (IDS) [26]. In continuous action space, the perturbation of the actions is provided by independent random noise, with (e.g., SAC [5]) or without (e.g., DDPG [4]) learned noise parameters. Unstructured exploration has several drawbacks [9], which have been tackled by ensuring a temporally coherent noise through colored noise [11], an auto-regressive process [27] or a Markov chain [28]. Another strategy to induce time correlation consists in keeping the perturbation parameters constant for multiple steps. This concept has been applied to perturbations of the network parameters [29] and to state-dependent action perturbations [9, 10, 8, 30]. While these exploration strategies provide smoother noise in time, Lattice aims to pair time-correlated state-dependent exploration with correlation in the action space.
Other exploration strategies. In on-policy RL, count-based exploration [32], curiosity-driven exploration [33] and random network distillation [34] are techniques to encourage visiting rare states. In off-policy RL, bootstrapped DQN [35] employs multiple Q-networks to guide exploration and improve performance. Additionally, unsupervised methods like Diversity Is All You Need [36] can achieve more diverse exploration by maximizing mutual information between a goal-conditioned policy and a set of skill labels. Starting to explore only after returning to a recently visited state allows to solve complex sequential-decision problem [37]. Lattice does not aim to replace curriculum learning [38-40, 15] or exploratory policies, but proposes exploring implicitly via the policy network in a monolithic way and can be combined with other approaches.
Exploration in over-actuated system. Overactuated systems, such as musculoskeletal systems, pose a challenge for exploration. Recent NeurIPS challenges have addressed the problem of learning control policies for locomotion and dexterous hand movements with musculoskeletal models [41, 14, 15]. The winners in these challenges have typically utilized complex training curricula, such as static-to-dynamic stabilization in the Baoding task, where a hand needs to accurately rotate two balls [15], or feature engineering and expert demonstration for locomotion [41]. Exploration can also be encouraged by learning a muscle-coordination network [42] or the state-dependent joint torque limits and energy function [43]. Similarly, Schumacher et al. proposed a self-organizing controller that outperforms previous approaches by inducing state-space covering exploration [19]. Integrating exploratory policies into a RL framework requires ad-hoc solutions, whereas Lattice can be trained end-to-end with both on-policy and off-policy algorithms.
3 Motivating example: The case of a flexor-extensor, single joint arm
To motivate Lattice, consider a simple flexor-extensor system, such as the elbow joint in the human upper arm (Fig. 2A). The activation  $a_{f} \in [0,1]$  of the flexor muscle (biceps) produces a negative angular acceleration at the elbow, while the activation  $a_{e} \in [0,1]$  of the extensor muscle (triceps) produces a positive angular acceleration. Using a first order approximation of the system dynamics, we can write the equation of the angular position of the elbow as  $\ddot{\theta} = \alpha (a_{e} - a_{r})$ , where  $\theta$  is the angle between the upper arm and the lower arm and  $\alpha$  [rad  $\mathrm{s}^{-2}$ ] is a parameter of the system which converts the muscle activations into angular acceleration. Given this system, we consider a simple control problem, where an agent needs to reach a target angular position of the elbow by activating the flexor and/or the extensor muscle. At every step, the agent observes the difference  $\Delta \theta = \theta_0 - \theta_t$  between the target angle  $\theta_0$  and the angle  $\theta_t$  at time  $t$ . For small values of  $\Delta \theta$ , we assume that the control policy learned by the agent can be approximated with the linear functions  $a_{e} = 0.5 + \Delta \theta$  and  $a_{f} = 0.5 - \Delta \theta$ . We now want to compare the effects on the angular acceleration produced by noise applied to the actions ( $a_{e}$  and  $a_{f}$ ) or to the latent state ( $\Delta \theta$ ) for various cases:
Case 1: action space noise. If we inject independent noise  $\epsilon_{f}$  and  $\epsilon_{e}$  following a normal distribution  $\mathcal{N}(0,\sigma^2)$  to the muscle activations, we have  $a_{f} \sim \mathcal{N}(\langle a_{f}\rangle ,\sigma^{2})$  and  $a_{e} \sim \mathcal{N}(\langle a_{e}\rangle ,\sigma^{2})$ . It can be shown that  $\ddot{\theta} \sim \mathcal{N}(\alpha (\langle a_e\rangle -\langle a_f\rangle),2\alpha^2\sigma^2)$  (Appendix A.1).
Case 2: latent space noise. If we instead inject Gaussian noise to the latent state  $\Delta \theta$ , while the marginal distribution of  $a_{f}$  and  $a_{e}$  is the same as with action space noise, the fact that the two distributions are correlated has an effect on the distribution of  $\ddot{\theta}$ . It can be shown that  $\ddot{\theta} \sim \mathcal{N}(\alpha (\langle a_{e} \rangle - \langle a_{r} \rangle), 4\alpha^{2}\sigma^{2})$  (Appendix A.1).
This toy example illustrates how perturbing the latent state, instead of the actions, leads to higher variance in the behavior space. In fact, when randomly perturbing two opposing muscles, half of the time the perturbations will be in opposite directions, with a reduced impact on the observed kinematics.
To corroborate this finding about the advantage of latent versus action perturbations, we considered a policy network trained to control a realistic musculoskeletal arm model, implemented in MyoSuite [18], featuring three flexor muscles and three extensor muscles. A neural network trained with PPO and Lattice exploration can learn to accurately reach a target angle in every episode, and while doing so learns to alternately activate the extensor or the flexor muscle group (Fig. 2 B, left). The training procedure is detailed in Appendix A.1. We collected a dataset of 100 episodes where the latent state of the policy network was perturbed with Gaussian noise. While applying independent random noise to the muscle activations produces a diagonal covariance matrix, perturbing the latent state of the network leads to a full covariance matrix (Fig. 2 B, right). Thus, as hypothesized, the last layer of the policy network learned to distinguish between agonist and antagonist muscles, namely, the covariance shows positive correlation among muscles of the same group. We then tested whether sampling action noise with this covariance matrix leads to higher variance in the task space (joint angles) than noise sampled independently for each muscle. We collected 100 episodes where, at each step, we computed an action with action space perturbation and one action with latent state perturbation. For consistency of evaluation, the perturbation in the action space was sampled with the same variance magnitude induced by the latent perturbation. We executed each pair of actions in two cloned environments and compared the variance of the joint angles. The latent state perturbation introduces higher variance in the joint angles (Fig. 2 C), thus driving more diverse kinematics and wider exploration. Overall, we conclude that latent exploration might be advantageous when learning to control over-actuated systems, and thus developed and empirically tested Lattice on different benchmarks.
4.1 LATent TIme-Correlated Exploration (Lattice)
In continuous control, the policy update of on-policy algorithms (e.g., REINFORCE [44], A3C [1], TRPO [21], PPO [2]) and of some off-policy algorithms (e.g., SAC [5]) requires the computation of  $\log \pi (a_t|s_t)$ , which is the logarithm of the probability of choosing a certain action  $a_{t}$  given the current state  $s_t$ , according to the current policy  $\pi$ . In such cases, the policy gradient needs to be
Algorithm 1 Standard (e.g., PPO, SAC)
Require: Policy  $\pi$  env. dynamics  $p$  std array  $\sigma^{(\mathbf{a})}$ $(\epsilon_{\mathbf{a}})_i\sim \mathcal{N}(0,\sigma_i^{(\mathbf{a})})$  1: Initialize state  $\mathbf{s}_t$    
2: step_count  $\leftarrow 0$    
3: while not done do   
4:  $\mathbf{x}_t\gets \pi (\mathbf{s}_t)$    
5:  $\epsilon_{\mathrm{a}}\gets$  sample  $(\sigma^{(\mathbf{a})})$    
6:  $\mathbf{a}_t\gets W\mathbf{x}_t + \epsilon_{\mathrm{a}}$    
7:  $\mathbf{s}_{t + 1},r_t$  done  $\leftarrow p(\mathbf{s}_t,\mathbf{a}_t)$    
8: step_count  $\leftarrow$  step_count  $+1$    
9: end while
Algorithm 2 Lattice
Require: Policy  $\pi$  ,env.dynamics  $p$  period  $T$  , matrix  $S^{(\mathbf{a})}$  , matrix  $S^{(\mathbf{x})}$ $\alpha \in \{0,1\}$ $(P_{\mathrm{a}})_{i,j}\sim \mathcal{N}(0,S_{i,j}^{\mathrm{(a)}})$ $(P_{\mathbf{x}})_{i,j}\sim \mathcal{N}(0,S_{i,j}^{(\mathbf{x})})$  1: Initialize state  $\mathbf{s}_t$  .   
2: step_count  $\leftarrow 0$  .   
3: while not done do   
4: if step_count mod  $T = 0$  then   
5:  $P_{\mathrm{a}}\gets$  sample  $(S^{(\mathbf{a})})$  .   
6:  $P_{\mathbf{x}}\gets$  sample  $(S^{(\mathbf{x})})$  .   
7: end if   
8:  $\mathbf{x}_t\gets \pi (\mathbf{s}_t)$  .   
9:  $\mathbf{a}_t\gets (W + P_\mathrm{a} + \alpha W P_\mathrm{x})\mathbf{x}_t$    
10:  $\mathbf{s}_{t + 1},r_t$  ,done  $\leftarrow p(\mathbf{s}_t,\mathbf{a}_t)$    
11: step_count  $\leftarrow$  step_count  $+1$  .   
12: end while
Algorithms 1-2: Experience collection algorithm with independent Gaussian noise (left) and with time-and actuator-correlated noise (right). In green, we highlight the parameters defining the independent noise. In brown, the elements inducing time-correlation (gSDE). In blue, those implementing actuator correlation through the perturbation of the latent state (Lattice). Note that for  $\alpha = 0$ , Lattice is equivalent to gSDE.
backpropagated through the density function of the action distribution. This is typically accomplished via the reparametrization trick, in which the policy network outputs the parameters of a differentiable probability density function, whose analytical expression is known, so that gradients can flow through the probability estimation. In the standard case, when we apply independent Gaussian noise to the action components, the action probability distribution can be parametrized as a multivariate Gaussian with diagonal covariance matrix (Algorithm 1). We propose that the weights of the policy network can be naturally used to parameterize the amount of exploration noise (Algorithm 2).
While a generic perturbation of the policy network's activations would lead to an action distribution with an unknown probability density function, in Lattice we limit the latent perturbation to the last layer's latent state, which is linearly transformed into the action (Fig. 1 A). Consider the output of the last layer of the policy network  $\mathbf{x} \in \mathbb{R}^{N_x}$  and the matrix  $W \in \mathbb{R}^{N_a \times N_x}$ , mapping the embedding state to an action according to  $\mathbf{a} = W\mathbf{x}$ . We have indicated the size of the latent space with  $N_x$  and the size of the action space with  $N_a$ . If we assume a perturbation of the latent state  $\mathbf{x}$  with independent Gaussian noise  $\epsilon = \mathcal{N}(\mathbf{0}, \Sigma_{\mathbf{x}})$ , where  $\Sigma_{\mathbf{x}}$  is a positive-definite diagonal matrix, then the perturbed latent state  $\tilde{\mathbf{x}}$  is distributed as  $\tilde{\mathbf{x}} \sim \mathcal{N}(\mathbf{x}, \Sigma_{\mathbf{x}})$ . Thus, the action distribution  $\mathbf{a} \sim \mathcal{N}(\bar{W}\mathbf{x}, W\Sigma_{\mathbf{x}}W^\top + \Sigma_{\mathbf{a}})$ , where with  $\Sigma_{\mathbf{a}}$  we have indicated the independent component of the action noise, can be derived as a linear transformation of a multivariate Gaussian distribution (details in Appendix A.2). This formula provides an analytical expression for  $\log \pi(\mathbf{a}_t|\mathbf{s}_t)$ , as a function of the policy network weights and the covariance matrix of the latent state:
$$
\log \pi (\mathbf {a} | \mathbf {s}) = - \frac {N _ {a}}{2} \log (2 \pi) - \frac {1}{2} \log | W \Sigma_ {\mathbf {x}} W ^ {\top} + \Sigma_ {\mathbf {a}} | - \frac {1}{2} (W \mathbf {x} - \mathbf {a}) ^ {\top} (W \Sigma_ {\mathbf {x}} W ^ {\top} + \Sigma_ {\mathbf {a}}) (W \mathbf {x} - \mathbf {a}). \tag {1}
$$
The probability of an action depends on the network weights both through its mean value  $W\mathbf{x}$  and through the perturbation variance matrices  $\Sigma_{\mathbf{x}}$  and  $\Sigma_{\mathbf{a}}$ . Depending on the RL algorithm, it might be convenient not to propagate the gradients of the policy network through the variance component of the loss. In this way, the expected action trains the policy network, while the loss due to the standard deviation can regulate its magnitude by updating the parameters  $\Sigma_{\mathbf{x}}$  and  $\Sigma_{\mathbf{a}}$ .
Figure 3: Learning curves in five locomotion tasks of PyBullet, representing the cumulative reward of each training episode (mean  $\pm$  s.e.m. across random seeds). Our baseline results improve over the benchmark rewards of the pre-trained networks of RL Zoo [46], possibly because we used 16 vectorized environments to collect transitions and therefore a higher number of steps. The period  $T$  refers to the time correlation of the exploration noise. Lattice, gSDE and Pink Noise perform comparably in all environments apart from the Humanoid [16], where Lattice is more efficient and achieves higher average reward.
4.2 Lattice generalizes time-correlated noise
Lattice can be thought of as an extension of Generalized State-Dependent Exploration (gSDE) [8], as it inherits all its properties, while extending it with the possibility of perturbing the latent state of the policy, besides the actions (Fig. 1 B). In Lattice, the agent learns two sets of parameters,  $S_{\mathbf{x}}$  and  $S_{\mathbf{a}}$ , representing the standard deviation of  $N_{\mathbf{x}} \times (N_{\mathbf{x}} + N_{\mathbf{a}})$  Gaussian distributions ( $N_{\mathbf{x}}$  and  $N_{\mathbf{a}}$  being the size of the latent space and of the action space, respectively). These parameters are used to sample, every  $T$  steps, two noise matrices  $P_{\mathbf{x}}$  and  $P_{\mathbf{a}}$  (Fig. 1 C). These matrices are used to determine the latent noise  $\epsilon_{\mathbf{x}} = \alpha P_{\mathbf{x}}\mathbf{x}$  and the action noise  $\epsilon_{\mathbf{a}} = P_{\mathbf{a}}\mathbf{x}$ . The parameter  $\alpha \in \{0,1\}$  can be used to turn on or off the correlated action noise. For  $\alpha = 0$ , Lattice is equivalent to gSDE. The value of  $T$  regulates the amount of time-correlation of the noise. If  $T$  is large,  $P_{\mathbf{x}}$  and  $P_{\mathbf{a}}$  are sampled infrequently, causing the noise applied to the latent state and to the action to be strongly time-correlated (Algorithm 2).
In short, the action distribution depends on the latent state  $\mathbf{x}$ , the last linear layer parameters  $W$ , the Lattice noise parameters  $S_{\mathbf{x}}$  and the gSDE noise parameters  $S_{\mathbf{a}}$  in the following way:
$$
\pi (\mathbf {a} | \mathbf {s}) = \mathcal {N} \left(W \mathbf {x}, \operatorname {D i a g} \left(S _ {\mathbf {a}} ^ {2} \mathbf {x} ^ {2}\right) + \alpha^ {2} W \operatorname {D i a g} \left(S _ {\mathbf {x}} ^ {2} \mathbf {x} ^ {2}\right) W ^ {\top}\right) \tag {2}
$$
4.3 Implementation details of Lattice
Code to study the implementation details is available at https://github.com/amathislab/lattice.
Noise magnitude. The magnitude of the variance of individual noise components can either be learned or kept fixed, and can or can not be made dependent on the current state. We extend the implementation of gSDE [8, 45], which makes the noise magnitude a learnable parameter independent of the policy network.
Controlling the variance. We introduce two clipping parameters to limit the minimum and the maximum value of the variance of each latent noise component, to avoid excessive perturbations and convergence to a deterministic policy. This clipping does not modify the analytical expression of the action distribution, because it does not affect the sampled values, but rather the variance of the distribution.