SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Layer-dependent noise rescaling. As the noise vectors are given by  $\epsilon_{\mathbf{x}} = \alpha P_{\mathbf{x}}\mathbf{x}$  and  $\epsilon_{\mathrm{a}} = P_{\mathrm{a}}\mathbf{x}$ , assuming the components of  $\mathbf{x}$  have similar magnitude, the noise scales with the size of the latent state. To remove this effect, we apply a correction to the learned matrices  $\log \tilde{S}_{\mathbf{x}}$  and  $\log \tilde{S}_{\mathbf{a}}$  before sampling  $P_{\mathbf{x}}$  and  $P_{\mathrm{a}}$ :  $\log S_{\mathbf{x}} = \log \tilde{S}_{\mathbf{x}} - 0.5\log (N_{\mathbf{x}})$  and  $\log S_{\mathbf{a}} = \log \tilde{S}_{\mathbf{a}} - 0.5\log (N_{\mathbf{x}})$ . It can be proven that this correction removes the dependence of the noise from the size of the latent state (Appendix A.3), and thus from the network architecture.
Figure 4: Learning curves in the MyoSuite environments, showing the adjusted solved fraction throughout the training (ratio between the number of steps in which the target is reached and the maximum length of the episode, mean  $\pm$  s.e.m. across random seeds). In the pose environments, where the target prescribes the angle of all joints, standard exploration can be better than Lattice. In the other environments, where the exploration can leverage the learned muscle synergies, Lattice is on par with or better than independent noise. We omit the learning curves with time correlation above 4, as they were always worse or much worse than 1 and 4 in PPO, in accordance with previous findings for locomotion [8].
Preventing singularity of the covariance matrix. The covariance matrix of the action distribution must be positive definite. It has the expression  $W\Sigma_{\mathbf{x}}W^{\top} + \Sigma_{\mathbf{a}}$  (Eq. 2), with  $\Sigma_{\mathbf{x}} = \alpha^{2}\mathrm{Diag}(S_{\mathbf{x}}^{2}\mathbf{x}^{2})$  and  $\Sigma_{\mathbf{a}} = \mathrm{Diag}(S_{\mathbf{a}}^{2}\mathbf{x}^{2})$ , where the square operations are to be intended element-wise. By construction, the covariance matrix is positive-semidefinite. However, it could be singular, e.g., when  $\mathbf{x}$  is the null vector (Appendix A.4 for further details). To prevent the singularity of the covariance matrix, we add a small positive regularization value to its diagonal terms.
We benchmarked Lattice on standard locomotion tasks [47, 6, 16, 48-50] in PyBullet [20], as well as musculoskeletal control tasks of MyoSuite [18] built in MuJoCo [31]. Both libraries include continuous control tasks of varying complexity. While in PyBullet the actuators apply a torque to each individual joint, in MyoSuite the agent controls its body through muscle activations. We give a complete description of each task in Appendix A.6. We implemented Lattice as an extension of gSDE in the RL library Stable Baselines 3 [45], which we used for all our experiments; the hyperparameters of the algorithms are detailed in Appendix A.7. All the results are averaged across 5 random seeds. The training was run on a GPU cluster, for a total of approximately 10,000 GPU-hours.
5.1 Pybullet locomotion environments
Lattice can be paired with off-policy RL algorithms, such as SAC [5]. We tested this combination in the locomotion environments of PyBullet, where gSDE-SAC achieves state-of-the-art performance [8, 46]. We also included Pink Noise exploration [11] as an additional baseline, which is the state-of-the-art colored noise process for exploration with off-policy RL algorithms. We used the same network architecture and hyperparameters for SAC specified in [46] for all the environments (see Appendix A.7). Preliminary experiments on the parameters of Lattice showed that environments with a smaller action space benefit from a higher initial standard deviation of the exploration matrix, so we set all the elements of  $\log S_{\mathbf{x}}$  and  $\log S_{\mathbf{a}}$  to 0 for Ant and Humanoid and to 1 for the other environments. We found that Lattice achieves similar performance to gSDE-SAC and SAC-pink for lower-dimensional morphologies (Ant, Hopper, Walker and Half Cheetah), while outperforming them substantially for the Humanoid (Fig. 3 and Appendix A.8). This suggests that Lattice can achieve
higher performance when controlling larger-dimensional morphologies. Furthermore, we later show that the policy learnt with Lattice-SAC for the Humanoid is more energy efficient than that learnt with SAC (Section 6).
5.2 Musculoskeletal control: MyoSuite environments
For musculoskeletal control, we tested several tasks from the MyoSuite [18]:
- Three pose tasks (Elbow Pose, Finger Pose and Hand Pose). In each session, a target angular position is sampled independently for each joint. This means that the policy has to learn how to control each degree of freedom independently.
- Two reach tasks (Finger Reach and Hand Reach). In each session, a target position for each finger tip (one for Finger Reach and five for Hand Reach) is sampled independently.
- Three object manipulation tasks (Reorient, Pend and Baoding). In each session, the target is a fixed or moving target for an object (a die in Reorient, a pen in Pen and two Baoding balls in Baoding).
We tested Lattice with a LSTM network and PPO. This choice is motivated by (our) winning solution of the Boading ball task in the 2022 NeurIPS MyoChallenge [15], which used PPO together with a LSTM network [51]. For our experiments we use the same PPO hyperparameters and network architecture, which we keep identical across methods (see Appendix A.7).
We set the initial value of all the elements of  $\log S_{\mathbf{x}}$  and  $\log S_{\mathbf{a}}$  to 0 in every task, except in Elbow pose and Hand pose, where we set it to 1.
In the reach and challenging object manipulation tasks (Reorient, Baoding and Pen), where the policy network has to control a complex hand model with 39 muscles and receives an observation with more than 100 dimensions, Lattice consistently outperforms the baselines (Fig. 4 and Appendix A.8). In contrast, in the pose tasks PPO performs comparably or better than Lattice-PPO (Fig. 4 and Appendix A.8). Those tasks define dense target states and, perhaps, exploration and coordination is less important (see below).
Interestingly, focusing action noise in the task-relevant space allows Lattice to avoid activating muscles unnecessarily, leading to conspicuous energy saving in the manipulation and reaching tasks (Fig. 5). In reach and object manipulation, Lattice-PPO achieves better reward at lower energy cost. We speculate that injecting correlated noise across muscle activations improves exploration in the
Figure 5: Energy saving versus reward gain for the muscle control tasks. Lattice learns policies  $20\%$  to  $60\%$  more energy efficient than independent exploration in the reaching and object manipulation tasks, while also achieving higher reward. In the pose tasks, independent exploration performs better, with similar energy consumption.
space of task-relevant body poses and facilitates the discovery of efficient, coordinated movements. We tested this hypothesis next.
6 How does Lattice explore?
In most of the considered musculoskeletal control environments, Lattice finds a more energy efficient control strategy, without compromising the performance (Fig. 5 and Appendix A.8). Especially in Reorient, where Lattice outperforms standard exploration in cumulative reward, it does so at a
Figure 6: A Left: Graph of the fraction of noise allocated to each group of action components by a stochastic policy trained in the Humanoid environment with SAC and Lattice-SAC. In the Humanoid, Lattice tends to focus most of the variance on the task-relevant actuators (legs). Middle: Distribution of the energy consumption. Each bar represents the number of test episodes falling in the corresponding energy consumption interval. Right: Cumulative explained variance of the actions' principal components, computed from a dataset of 100 test episodes per seed. Shaded area represents  $95\%$  confidence interval across training seeds. More principal components are required to explain the same fraction of variance in SAC and PPO versus Lattice-SAC and Lattice-PPO. B Heatmap of the correlation matrix for the action space (left column) and noise of the action (right column) for SAC (top row) and Lattice-SAC (bottom row). C-D Same as A-B but analyzing a policy trained on the MyoSuite Reorient task with PPO and Lattice-PPO. Lattice tends to induce sharper correlation between action components. Local patterns of correlation between action components can be recognized in the noise covariance matrix.
fraction of the energy cost. We hypothesize that this is a direct consequence of biasing the noise with the same correlation as the actuators, preventing energy-consuming co-activations, which the agent needs to compensate for, despite the reward not being affected by them.
We assessed whether coordination emerged by analyzing policies for the Humanoid locomotion task and the Reorient task. By looking at the relative allocation of action noise across actuators (Fig. 6 A, C), we can see that in the Humanoid, compared to SAC, Lattice-SAC re-directs more exploration noise towards the task-relevant leg motors (SAC:  $40\%$  Legs,  $45\%$  Arms and  $15\%$  Body; Lattice-SAC:  $50\%$  Legs,  $32\%$  Arms,  $18\%$  Body). In Reorient, where all the muscles contribute to controlling the hand, the difference is less evident. Furthermore, both in the Humanoid and Reorient task, the intrinsic dimensionality of the actions output by a policy trained with Lattice is lower. Indeed, consistently across random seeds, fewer principal components explain a higher fraction of the action variance in Lattice than in standard exploration (Fig. 6 A, C). Lattice consistently promotes policies with actions lying on a lower-dimensional manifold, and this effect extends beyond Humanoid and Reorient, to almost all the considered tasks (Appendix A.9). We speculate that this property of Lattice determines when it is to be preferred to uncorrelated noise, i.e., when there exists a low-dimensional solution to the task which allows the agent to take advantage of motor synergies. In environments where the dimensionality of the task is intrinsically low (reach and object manipulation), Lattice achieves strong performance. Further experiments in the pose and in the object manipulation tasks, performed with Lattice-SAC, confirm the result that Lattice-SAC is better than SAC for learning to manipulate objects (Appendix A.8).
The intrinsic dimensionality of the actions decreases as an effect of training (Appendix A.9). While training tends to reduce the dimensionality of the policy regardless of the action noise, the decrease is more marked with Lattice. The low dimensionality of the actions is explained by the action correlation matrices, which show increased cross-actuator coordination with Lattice (Fig. 6 B, D). We speculate that this is driven by the correlation structure of the noise. While the off-diagonal elements are close to 0 with uncorrelated noise, those of Lattice present a structure that resembles that of the action correlation matrix (Fig. 6 B, D), also consistent with the motivating example (Section 3). Indeed, if we consider the off-diagonal elements at position  $(i,j)$ , we have that  $\mathrm{Cov}(a_i,a_j) = W\mathrm{Cov}(\mathbf{x})W^\top$  while  $\mathbb{V}[\pi (\mathbf{a}|\mathbf{s})]_{i,j} = \alpha^2\left(W\mathrm{Diag}(S_{\mathbf{x}}^2\mathbf{x}^2)W^\top\right)_{i,j}$  (further details in Appendix A.5). In the case where  $\alpha = 1$ ,  $S_{\mathbf{x}}$  is the identity and the components of  $\mathbf{x}$  are uncorrelated, then the two matrices are identical. While throughout the training the parameters of  $S_{\mathbf{x}}$  can adapt, we can empirically see that the covariance matrix retains elements of its original structure.
We proposed Lattice, a method that leverages the synergies between action components encoded in the weights of the policy network to drive exploration based on the latent state. Latent exploration has proven effective in discovering high-reward and energy efficient policies in complex locomotion and muscle control tasks. Remarkably, Lattice outperforms random exploration in the reach and object manipulation tasks, with a large reduction in energy consumption. We showed that perturbing the last layer of the policy network introduces enough bias into the action noise. Importantly, this specific perturbation allowed us to find an analytical expression for the action distribution, which makes Lattice a straightforward enhancement for any on-policy or off-policy continuous control algorithm.
This sets Lattice apart from other forms of exploration using an auxiliary policy to collect experience, as they see their application limited to off-policy RL. Furthermore, knowing the action distribution is fundamental for those off-policy algorithms learning a stochastic policy, such as SAC. On the other hand, modeling the action distribution with a multivariate Gaussian with full covariance matrix comes at additional computational cost. In fact, Lattice introduces a training overhead over gSDE and PPO/SAC (approx.  $20\% - 30\%$ , depending on the hardware). This is due to the additional matrix multiplications required to estimate the action probabilities (Eq. 1). This limitation opens an interesting research direction, e.g., introducing sparsity constraints in the distribution matrices. For example,  $S_{\mathrm{a}}$  and  $S_{\mathrm{x}}$  could be forced to be diagonal or low-rank, without preventing motor synergies to be included into the noise distribution. We leave this investigation to future work.
Biological motor learning finds efficient and robust solutions [12, 13, 52-55]. Lattice also discovers low-energy solutions and it will be an interesting future question, if the brain is also performing some form of latent-driven exploration. Beyond this specific question, we are enthusiastic that advances in musculoskeletal simulators [18] as well as reinforcement learning are opening up exciting avenues for Neuroscience [56].
Our research aims to facilitate the training of motor control policies for complex tasks. Lattice might be employed as a tool in Artificial Intelligence, Robotics and Neuroscience, contributing to the creation of powerful autonomous agents. The energy efficiency of the learned policies might be of great interest for energy-sensitive robotics applications for reducing the carbon footprint of humans. While the advancement of the research in this field is a great opportunity, the related concerns should be carefully addressed [57].
We are grateful to Adriana Rotondo, Stephane D'Ascoli and other members of the Mathis Group for comments on earlier versions of this manuscript. Our work was funded by Swiss SNF grant (310030_212516) and EPFL. A.M.V.: Swiss Government Excellence Scholarship. A.Z.H.: EPFL School of Life Sciences Summer Research Program. We thank the Neuro-X institute at EPFL for supporting the travel expenses (A.S.C. and A.M.V).
[54] Dagmar Sternad. It's not (only) the mean that matters: variability, noise and exploration in skill learning. Current opinion in behavioral sciences, 20:183-195, 2018.
[55] Adam S Lowet, Qiao Zheng, Sara Matias, Jan Drugowitsch, and Naoshige Uchida. Distributional reinforcement learning in the brain. Trends in neurosciences, 43(12):980-997, 2020.
[56] Sébastien B Hausmann, Alessandro Marin Vargas, Alexander Mathis, and Mackenzie W Mathis. Measuring and modeling the motor system with machine learning. Current opinion in neurobiology, 70:11-23, 2021.
[57] Maartje MA de Graaf. An ethical evaluation of human-robot relationships. International journal of social robotics, 8:589–598, 2016.
[58] Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pages 249–256. JMLR Workshop and Conference Proceedings, 2010.
A.1 Singe-joint arm: detailed calculations 15
A.2 Lattice's Action Distribution Parameterization 16
A.3 Noise rescaling across networks of different size 17
A.4 Conditions on the covariance matrix of the action distribution 17
A.5 Empirical covariance of the action 18
A.6 Parameters of the PyBullet and MyoSuite environments 18
A.7 Hyperparameters of SAC, PPO, gSDE and Lattice 18
A.8 Detailed reward and energy results of all experiments 19
A.9 Evolution of the dimensionality of the policy during training 19
A.1 Singe-joint arm: detailed calculations
Case 1: action space noise. By combining the equation of the angular acceleration  $\ddot{\theta} = \alpha (a_{e} - a_{f})$  with the distributions of the muscle activations  $a_{f}\sim \mathcal{N}(\langle a_{f}\rangle ,\sigma^{2})$  and  $a_{e}\sim \mathcal{N}(\langle a_{e}\rangle ,\sigma^{2})$ , we can compute expected value and variance of the angular acceleration:
$$
\mathbb {E} \left[ \ddot {\theta} \right] = \mathbb {E} \left[ \alpha \left(a _ {e} - a _ {r}\right) \right] = \alpha \left(\mathbb {E} \left[ a _ {e} \right] - \mathbb {E} \left[ a _ {f} \right]\right) = \alpha \left(\left\langle a _ {e} \right\rangle - \left\langle a _ {r} \right\rangle\right) \tag {3}
$$
$$
\mathbb {V} \left[ \ddot {\theta} \right] = \mathbb {V} \left[ \alpha \left(a _ {e} - a _ {r}\right) \right] = \alpha^ {2} \left(\mathbb {V} \left[ a _ {e} \right] + \mathbb {V} \left[ a _ {f} \right]\right) = 2 \alpha^ {2} \sigma^ {2} \tag {4}
$$
In the second derivation we have used the fact that  $a_{e}$  and  $a_{r}$  are independent random variables to compute their variances separately, and that  $\mathbb{V}[-a_f] = \mathbb{V}[a_f]$ .
Case 2: latent space noise. In this case we apply noise to the latent state, so that  $\Delta \theta \sim \mathcal{N}(\langle \theta \rangle ,\sigma^2)$ . While it still holds that  $a_{f}\sim \mathcal{N}(\langle a_{f}\rangle ,\sigma^{2})$  and  $a_{e}\sim \mathcal{N}(\langle a_{e}\rangle ,\sigma^{2})$ , the two random variables are no longer independent. This fact does not change the computation of the expectation:
$$
\mathbb {E} \left[ \ddot {\theta} \right] = \mathbb {E} \left[ \alpha \left(a _ {e} - a _ {r}\right) \right] = \alpha \left(\mathbb {E} \left[ a _ {e} \right] - \mathbb {E} \left[ a _ {f} \right]\right) = \alpha \left(\left\langle a _ {e} \right\rangle - \left\langle a _ {r} \right\rangle\right) \tag {5}
$$
For the variance, we need to make the dependence on  $\theta$  explicit, using the formulas  $a_{e} = 0.5 + \Delta \theta$  and  $a_{f} = 0.5 - \Delta \theta$ :
$$
\mathbb {V} \left[ \ddot {\theta} \right] = \mathbb {V} \left[ \alpha \left(a _ {e} - a _ {r}\right) \right] = \mathbb {V} \left[ 2 \alpha \Delta \theta \right] = 4 \alpha^ {2} (\mathbb {V} [ \Delta \theta ]) = 4 \alpha^ {2} \sigma^ {2} \tag {6}
$$
We can observe that the variance is double in the case of latent space noise.
To test whether these results can be observed also in simulation, we considered the Elbow pose task of MyoSuite, where we trained a policy with PPO and Lattice exploration (with the hyperparameters specified in Table T4). The policy reaches a solved value of 0.95, meaning that the angle between the upper and the lower arm is in the immediate neighborhood of the target angle  $95\%$  of the simulation time. We assessed the effect of latent and action perturbations on the angular position by carrying out two simultaneous simulations of the environment per noise type. At every step, one simulation received a noisy action, while the other would receive a deterministic, noise-free action. We then registered the difference in elbow angle between the noisy and the deterministic simulation. To avoid having the results obfuscated by the cumulative perturbation over longer time periods, we synchronized the simulators after every step at the state achieved through the stochastic simulation.
To compute comparable results with latent and action noise, we had to make sure that the variance of each action component would be equivalent in both cases. We accomplished this by first running the simulation with latent noise, with the same variance for each latent state component. We then measured the variance that such perturbation would cause to each action component, and we used
this value to generate the action perturbation noise. This ensures that the difference in the variability of the elbow angle when applying latent or action noise is not due to a difference in scale, but only to the presence of off-diagonal elements in the noise covariance matrix.
A.2 Lattice's Action Distribution Parameterization
In Lattice, the action vector  $\mathbf{a}$  is a linear transformation of the perturbed latent vector  $\tilde{\mathbf{x}} = \mathbf{x} + \epsilon_{\mathbf{x}}$  further perturbed with independent action noise  $\epsilon_{\mathbf{a}}$ :
$$
\mathbf {a} = W \tilde {\mathbf {x}} + \epsilon_ {\mathbf {a}} = W \left(\mathbf {x} + \epsilon_ {\mathbf {x}}\right) + \epsilon_ {\mathbf {a}} \tag {7}
$$
Furthermore, the latent noise  $\epsilon_{\mathbf{x}}$  and the action noise  $\epsilon_{\mathbf{a}}$  are defined as follows:
$$
\epsilon_ {\mathbf {x}} = P _ {\mathbf {x}} \mathbf {x} \quad \text {a n d} \quad \epsilon_ {\mathbf {a}} = P _ {\mathbf {a}} \mathbf {x} \tag {8}
$$
The elements of  $P_{\mathbf{x}}$  and  $P_{\mathbf{a}}$  are distributed as independent Gaussians:
$$
\left(P _ {\mathbf {x}}\right) _ {i, j} \sim \mathcal {N} \left(0, \left(S _ {\mathbf {x}}\right) _ {i, j} ^ {2}\right) \quad \text {a n d} \quad \left(P _ {\mathbf {a}}\right) _ {i, j} \sim \mathcal {N} \left(0, \left(S _ {\mathbf {a}}\right) _ {i, j} ^ {2}\right) \tag {9}
$$
where  $S^{(\mathbf{x})}$  and  $S^{(\mathbf{a})}$  are learnt parameter matrices whose elements represent the standard deviation of each element of the perturbation matrices. We can therefore compute the distribution of each element of the noise vectors. They are defined as:
$$
\left(\boldsymbol {\epsilon} _ {\mathbf {x}}\right) _ {i} = \alpha \sum_ {j = 1} ^ {N _ {\mathbf {x}}} \left(P _ {\mathbf {x}}\right) _ {i, j} \mathbf {x} _ {j} \quad \text {a n d} \quad \left(\boldsymbol {\epsilon} _ {\mathbf {a}}\right) _ {i} = \sum_ {j = 1} ^ {N _ {\mathbf {x}}} \left(P _ {\mathbf {a}}\right) _ {i, j} \mathbf {x} _ {j} \tag {10}
$$
meaning that they are the sum of independent Gaussian random variables. We can compute their mean and variance as follows:
$$
\mathbb {E} \left[ \left(\boldsymbol {\epsilon} _ {\mathbf {x}}\right) _ {i} \right] = \mathbb {E} \left[ \alpha \sum_ {j = 1} ^ {N _ {\mathbf {x}}} \left(P _ {\mathbf {x}}\right) _ {i, j} \mathbf {x} _ {j} \right] = \alpha \sum_ {j = 1} ^ {N _ {\mathbf {x}}} \mathbf {x} _ {j} \mathbb {E} \left[ \left(P _ {\mathbf {x}}\right) _ {i, j} \right] = 0 \tag {11}
$$
$$
\mathbb {V} \left[ \left(\boldsymbol {\epsilon} _ {\mathbf {x}}\right) _ {i} \right] = \mathbb {V} \left[ \alpha \sum_ {j = 1} ^ {N _ {\mathbf {x}}} \left(P _ {\mathbf {x}}\right) _ {i, j} \mathbf {x} _ {j} \right] = \alpha^ {2} \sum_ {j = 1} ^ {N _ {\mathbf {x}}} \mathbf {x} _ {j} ^ {2} \mathbb {V} \left[ \left(P _ {\mathbf {x}}\right) _ {i, j} \right] = \alpha^ {2} \sum_ {j = 1} ^ {N _ {\mathbf {x}}} \mathbf {x} _ {j} ^ {2} \left(S _ {\mathbf {x}}\right) _ {i, j} ^ {2} \tag {12}
$$
$$
\mathbb {E} \left[ \left(\boldsymbol {\epsilon} _ {\mathbf {a}}\right) _ {i} \right] = \mathbb {E} \left[ \sum_ {j = 1} ^ {N _ {\mathbf {x}}} \left(P _ {\mathbf {a}}\right) _ {i, j} \mathbf {x} _ {j} \right] = \sum_ {j = 1} ^ {N _ {\mathbf {x}}} \mathbf {x} _ {j} \mathbb {E} \left[ \left(P _ {\mathbf {a}}\right) _ {i, j} \right] = 0 \tag {13}
$$
$$
\mathbb {V} \left[ \left(\boldsymbol {\epsilon} _ {\mathbf {a}}\right) _ {i} \right] = \mathbb {V} \left[ \sum_ {j = 1} ^ {N _ {\mathbf {x}}} \left(P _ {\mathbf {a}}\right) _ {i, j} \mathbf {x} _ {j} \right] = \sum_ {j = 1} ^ {N _ {\mathbf {x}}} \mathbf {x} _ {j} ^ {2} \mathbb {V} \left[ \left(P _ {\mathbf {a}}\right) _ {i, j} \right] = \sum_ {j = 1} ^ {N _ {\mathbf {x}}} \mathbf {x} _ {j} ^ {2} \left(S _ {\mathbf {a}}\right) _ {i, j} ^ {2} \tag {14}
$$
The covariance of noise elements at different indices is 0. Indeed, for  $i \neq j$ :