Masked Sensory-Temporal Attention for Sensor Generalization in Quadruped Locomotion
Dikai Liu Tianwei Zhang Jianxiong Yin Simon See

Rethinking Improved Privacy-Utility Trade-off with Pre-existing  Knowledge for DP Training
    
        Yu Zheng, 
        Wenchao Zhang, 
        Yonggang Zhang^∗, 
        Wei Song,  
        Kai Zhou, 
        Bo Han 
    
    ∗ Yonggang Zhang is the corresponding author. 
    Yonggang Zhang is with the Department of Computer Science, University of Technology Sydney, Australia. E-mails: yonggang.zhang-1@uts.edu.auYu Zheng is with the Department of Information Engineering, Chinese University of Hong Kong, Shatin, Hong Kong SAR. E-mail: yuzheng404@link.cuhk.edu.hkWenchao Zhang is with the College of Computer Science and Technology, Xi’an University of Science and Technology, China. E-mail: zhangwenchao@xust.edu.cn; Wei Song is with the School of Computer Science and Engineering, Northeastern University, China.  E-mail: songwei@mail.neu.edu.cn; Kai Zhou is with the Department of Computing, Hong Kong Polytechnic University, Kowloon, China. E-mail: kaizhou@polyu.edu.hk; Bo Han is with the Department of Computer Science, Hong Kong Baptist University, Hong Kong SAR. E-mails: bhanml@comp.hkbu.edu.hk
    

    MOX - Dipartimento di Matematica “F. Brioschi”, Politecnico di
Milano, via Bonardi 9, 20133 Milan, Italy

    ^1 arashandrea.roknian@mail.polimi.it

    ^2 anna.scotti@polimi.it

    ^3 alessio.fumagalli@polimi.it
    September 9, 2024
==================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================================



[1]NVIDIA AI Technology Centre (NVAITC); e-mail: {dikail,jianxiongy,ssee}@nvidia.com
[2]College of Computing and Data Science, Nanyang Technological University, Singapore; e-mail: dikai001@e.ntu.edu.sg, tianwei.zhang@ntu.edu.sg
[3]also with Nanyang Technological University and Coventry University



§ ABSTRACT 
    With the rising focus on quadrupeds, a generalized policy capable of handling different robot models and sensory inputs will be highly beneficial. Although several methods have been proposed to address different morphologies, it remains a challenge for learning-based policies to manage various combinations of proprioceptive information. This paper presents  (), a novel transformer-based model with masking for quadruped locomotion. It employs direct sensor-level attention to enhance sensory-temporal understanding and handle different combinations of sensor data, serving as a foundation for incorporating unseen information. This model can effectively understand its states even with a large portion of missing information, and is flexible enough to be deployed on a physical system despite the long input sequence.



 





§ INTRODUCTION


Quadrupedal robots have showcased their capability to navigate in various challenging terrains, thanks to the rapid advancements of deep reinforcement learning (RL) technology <cit.>. With the increasing availability of affordable quadruped robots on the market,

there is a growing interest in developing a general-purpose locomotion policy that can fit all types of quadrupedal devices. Typically learned locomotion policies are specific to the robot, observation space, and task they were trained on, making it challenging to transfer or generalize to other robots or scenarios.

Recent works, such as GenLoco <cit.> and ManyQuadrupeds <cit.>, have successfully developed a single policy with the ability to adapt to diverse morphologies and degrees of freedom (DoF). However, these models depend on a fixed observation space input to generate latent space representations. They become ineffective when facing the following situations: (1) deployed to quadrupeds equipped with a different sensor set, (2) sensor data turn to unreliable due to wear and tear, (3) new input is needed for new task. Since sensory feedback is interrelated and each sensor plays a critical role at different stages of locomotion <cit.>, a policy with a deep understanding of proprioceptive information to handle flexible input is desired, to enhance the generalization, flexibility, and extensibility.


















One promising solution is self-attention-based transformers <cit.>, which have demonstrated exceptional capabilities in understanding complex sequential information of arbitrary lengths. 

They have been widely used in robotics to enhance various tasks

with multimodal processes <cit.>. However, due to the complex model structures and vast parameters, robots driven by transformers often run at very low frequency <cit.>, or depend on external high-power computing platforms <cit.>. 

To address these issues, two common strategies are proposed to achieve smooth and stable control for transfermer-based locomotion tasks. The first strategy involves outputting high-level commands as an interface, such as gait pattern <cit.> and velocity commands <cit.>. An external joint controller, which requires additional design and training, is then used to convert these high-level commands into joint commands. However, the application of this method is restricted to the policy combined with perception and human-robot interaction.  The second strategy, targeting the vanilla locomotion tasks, applies linear projections to observation-action data for tokenization <cit.>. Although this method is straightforward and efficient for producing joint commands in an end-to-end manner, it limits the transformer's direct access to sensory information and still relies on fixed input, thereby constraining its in-context understanding capability and multimodal nature of the data.

To address the above limitations, we propose , a novel transformer-based model for end-to-end quadruped locomotion control. It achieves sensor input generalization with its multimodal nature, while still being directly deployable on physical systems. Specifically in , all sensory data are discretized and tokenized to form a long proprioceptive information sequence. Inspired by the work <cit.> on learning spatiotemporal information in video understanding, a random mask is applied to remove a portion of the observation during training. This significantly enhances the model's sensory-temporal understanding, to better handle different combinations of sensor data and serve as a foundation for incorporating unseen data. Additionally, it aids in identifying the most essential sensory information, thereby reducing the computational power required for physical deployment. Extensive experiments conducted in simulation demonstrate that our method can efficiently handle incomplete sensory information, even with half of the data missing. It is also robust against unseen data, making it a solid foundation for further extensions. With direct sensory-temporal attention, the model is flexible enough to mix-and-match desired information for finetuneing, meeting the requirement for different end-to-end quadruped locomotion control in the physical world. 




§ RELATED WORK




 §.§ Reinforcement Learning in Legged Locomotion

Reinforcement learning (RL) has gained significant attention in developing robotic controllers for tasks such as legged locomotion <cit.>, eliminating the need for extensive prior knowledge. With the advance of robotic simulation, RL-based locomotion is often trained in virtual environments <cit.> with diverse terrains <cit.> and randomized environmental factors <cit.> to improve the robustness of the agent. This technique is commonly known as domain randomization (DR). Training in simulators also provides rich information, some of which is not easily accessible in the real world (i.e., privileged information). To better interpret such information and bridge the sim-to-real gap when deploying policies to the physical world, system identification is commonly used to transfer knowledge to a deployable student policy. <cit.> employed action imitation to infer teacher behaviors using historical proprioceptive data. <cit.> further developed a two-stage adaption framework for faster and more robust online transfer, which has become the foundation for many subsequent works <cit.>. Another approach combines transfer loss with RL loss for joint optimization <cit.> to allow the student to explore freely to maximize the reward return.



 §.§ Transformer in Robotics


New models and methods are proposed to leverage the power of transformers in RL and robotics. For instance, Decision Transformer <cit.> converts states, actions, and rewards into embeddings using an encoder. Trajectory Transformer <cit.> uses the complete discretized trajectory for language model-like autoregressive prediction. Building on these frameworks, Gato <cit.> was developed to serve as a generalist agent for hundreds of tasks, including real-world robotic manipulation. Vision Language Models (VLM) <cit.> and Vision Language Action Models (VLA) <cit.> use transformers as interfaces for scene and language understanding to provide high-level commands for robotic control and human-robot interaction. 

For legged locomotion, <cit.> developed a transformer model for vision-based locomotion. It outputs high-level velocity commands and relies on a dedicated low-level controller for motor control. To bring transformers to direct motor control, <cit.> proposed TERT, which utilizes historical observation-action pairs to generate target motor commands directly. <cit.> applied a similar method to the bipedal locomotion task and later reformulated it as a next token prediction problem <cit.>.



§ PRELIMINARY

We adopt the two-stage teacher-student transfer approach from TERT <cit.> as the basis of our framework. It utilizes a well-trained teacher policy through RL with privileged information.



 §.§ Simulation Environment


We implement the simulation environment based on Isaac Gym and its open source library IsaacGymEnvs <cit.> to enable massive parallel training. 


Terrain and Curriculum. We adopt the terrain curriculum from <cit.> to achieve locomotion in dynamic terrains. Five terrain types (smooth slope, rough slope, stairs up, stairs down, discrete obstacle) are used to spawn a total of 20 terrains with proportions of [0.1, 0.1, 0.35, 0.25, 0.2] respectively. Each terrain has 10 levels with the increasing difficulty for larger slope angle, higher stairs step and larger obstacle size. During training, the agent progresses to the next level when the tracked linear reward reaches 80% of the maximum achievable value and regresses if it fails to reach 25%. If any agent finishes the highest level, it will be sent to a random level to continue exploring. 

Domain Randomization.

To enhance the robustness of the policy, DR is used in the simulation. The commanded longitudinal and lateral velocity is sampled in [-1.0, 1.0] m/s and horizontal angular velocity in [-1.0, 1.0] rad/s. The command is resampled every 5 seconds without any curriculum. Due to the significant computational power and time required for transformer calculations, system delay is introduced to simulate the delay in action updates.

Please refer to the Appendix for the complete DR parameters.





















Observations and Actions. The privilege observation e_t for teacher training contains ground truth data gathered fro simulation, including base linear and angular velocity, orientation, surrounding height map and randomized parameters as described above. For proprioceptive information, we use three commonly seen low-level sensors from quadrupeds, joint encoders, IMU and foot contact sensors. These sensors can provide five sensory data, including joint position q ∈ℝ^12, joint velocity q̇∈ℝ^12, angular velocity ω∈ℝ^3, gravity vector g ∈ℝ^3 and binary foot contact c ∈ℝ^4. Furthermore, the randomly sampled user command target cmd = [v_x, v_y, ω_z] and actions from previous step a_t-1∈ℝ^12 is added, resulting in a observation of o_t ∈ℝ^49 for each step. To gather the temporal information, a list of historical proprioceptive information [o_t, o_t-1, ⋯, o_t-T] from past T=15 steps is stored and passed to the student network. Thus, full observation is in the ℝ^49 × 15 space. Both teacher and student outputs the desired joint position a_t, which is further processed by a PD controller for the output torque τ = K_p (q̂-q) + K_d (q̂̇̂- q̇), with base stiffness and damping set to 30 and 0.7 respectively and the target joint velocity q̂̇̂ set to 0.

Reward Function for RL. The reward functions are designed to encourage the agent to follow the commanded velocity. Following <cit.>, we primarily penalize linear and angular movement along  other axes, large joint acceleration and excessive power consumption.

Please refer to the Appendix for the complete reward structure. 
























 §.§ Teacher Policy and Training


We implement a teacher policy following <cit.>. The teacher first encodes the privilege information e_t, with a factor encoder μ into a latent space l_t, which is then combined with the latest observation-action pair o_t for the teacher policy π̂ to output the desired joint position â_t:

    l̂_̂t̂ = μ(e_t),    â_t = π̂ (l̂_̂t̂, o_t)


The μ and π networks are implemented as MLP with hidden layers of [512, 256, 128] and [256, 128], respectively. The teacher policy is trained with PPO <cit.> directly to maximize the reward return and is shared across all student transfers at later stages for a fair comparison.



§ METHODOLOGY


We present , a novel transformer-based framework to generate a generalized understanding of low-level proprioceptive information for quadruped locomotion in complex terrains. With this foundational understanding, this framework is capable of handling different combinations of sensory inputs, enabling better generalization and flexibility. It can potentially be extended to incorporate high-dimensional sensors for more complex tasks. Figure <ref> shows the overview of .









Sensory Data Tokenize. To avoid using linear projection like previous transformer-based locomotion controllers <cit.>, we need to map continuous data to tokens directly. Following previous works <cit.>, we first normalize the sensory data based on the collected mean and variance, and then pass them through an encoder. This is similar to Gato <cit.>, which discretizes the value into 256 bins. The data are further mapped into a learnable embedding space with 128 dimensions.




Positional Embedding and Sensor Type Embedding. To distinguish proprioceptive information from different sensors with multiple dimensions and historical relations, two additional embeddings are added. The first one is a fixed 2D sin-cos position embedding <cit.> applied on the sensor dimension and time axis. This allows the framework to handle sensors with varying lengths of dimensions and historical time windows and be easily extendable. To accommodate the multimodal nature of the sensory data, another learnable embedding is add to indicate each sensor type. This enables easy mix and match of information from different sensors without concerns about the order or placeholders.



Random Masking. Inspired by the use of masking in image and video understanding <cit.>, we randomly mask out portions of the sensory information. This brings two benefits. First, since only part of the sensory data are visible to the network, the model is required to infer and reconstruct the missing information from them, thereby enhancing its understanding of the relationships between different sensory inputs. Second, random masking significantly reduces the training time and computational resources required, as the complexity of self attention is necessarily quadratic in the input length, making it more feasible to run in massive parallel.

Transformer Model. We implement a transformer model to process the generated tokens. The model consists of multi-head self-attention blocks with a MLP ratio of 2.0. An additional learnable state embedding  is added to consolidate the processed information <cit.>, which is subsequently projected into the action space with a MLP layer π:

    l_t = ([o_t, o_t-1, ⋯, o_t-T]),      a_t = π (l_t)



Teacher-Student Transfer. Following TERT <cit.>, we train  with a two-stage transfer strategy. In the first offline pretraining stage, trajectory is gathered by unrolling the well-trained teacher policy while the student will predict the next actions. This is to ensure that the student can produce reasonable actions during the second online correction stage to overcome the gap of distribution shift by training on its own trajectory. In each stage, 50 million trajectories are collected to train the transformer for 80k updates. 

We minimize the loss for the next step prediction:

    ℒ =  a_t - â_t  ^2






§ EXPERIMENTS AND RESULTS


We design and conduct various simulation experiments to evaluate the effectiveness of the proposed masking mechanism, overall performance of , and its generalization ability for different sensor data. We mainly focus on three metrics: linear velocity tracking return per step, angular velocity tracking return per step, and the total final reward return, which tells us how the agent can conduct the task (following users' commands) and what the overall performance is. All reported results are averaged over 5000 trails with five terrain types and different levels. They are normalized on the basis of respect teacher data for easy comparison. 



 §.§ Impact of Mask Ratio


First, we would like to investigate the maximum portion of missing data  can handle to reconstruct robot states. During the transfer stages, we set the masking ratio to 0%, 25%, 50% and 75% independently and test the 16 model with these masking ratios. Figure <ref> shows the resultant heatmap matrix. When trained without masking, despite the model having very good performance when all information is present, it suffers from missing data and cannot efficiently reconstruct the status. Also, we can see that the performance is more dependent on the masking ratio in the second stage than that in the first stage. This is because in the second transfer stage, the student is actually interacting with the environment to reduce the gap caused by missing information and observation shift. In contrast, the mission of the first stage is to generate a usable policy that outputs reasonable actions so the agent does not fail dramatically and has the chance in the second stage to generate high-quality trajectories for optimization, which is achievable even with a masking ratio of 75%. This also demonstrates the importance of using two-stage transfer. 

Comparing the performance of these models, we choose the one trained with the masking ratio of 75% in the first stage and 50% in the second stage, which can well balance the resource usage and performance. This selected model can restore almost full performance even under a masking ratio of 50% during testing. Such ratio can significantly reduce the training time and resources required. 





 §.§ Comparison with Baselines


We compare  with two baselines. The first is RMA <cit.>, which is implemented with TCN <cit.> to capture temporal information. The second is TERT <cit.>, a transformer-based framework with linear projection for tokenization of observations and actions in two favors: concatenated single token and separate tokens for states and action, resulting in T and 2T tokens respectively <cit.>. Both RMA and TERT have access and need the full observation information. To further evaluate the importance and capability of the transformer structure in our method, we replace it with a GRU <cit.> model. For both  and GRU, we apply a testing mask of up to 50%, as identified in Section <ref>. All baselines are trained with the same two-stage transfer method sharing a common well-trained teacher network. 
Table <ref> shows the comparison results. When fully optimized with the teacher-student transfer, the performance of all trained policies with complete observations is very close, often within just 2% difference. When faced with incomplete information, transformer-based  can have a better understanding of the data and reconstruct the robot state more accurately than the GRU-based network, even with only half of the information.  





 §.§ Generalization, Robustness and Flexibility


While achieving state-of-the-art performance,  additionally offers the benefit of generalization and flexibility to customize the model after training or even on the fly to fit the deployment requirement. We can balance the performance and computation power required by randomly masking out the observation directly, just like we do during testing. As quadrupedal robots are often not equipped with high-power computers and need to be off the grid to work remotely, such balance is crucial. The in-context understanding capability of the transformer can give us more insight. 



Important Sensory Feedback. As our random masking is applied to each dimension independently, even with a mask of 50%, the proprioceptive information is often not completely wiped out and the transformer can still trace it in the given time window. We further investigate the impact of removing one sensor completely from the feedback. The results are shown in Figure <ref>. It is clear that certain feedback like q̇, c, ω and even a_t-1 are quite redundant and a well trained transformer-based  can easily compensate the missing information.

















Missing of Sensor Dimension. For some proprioceptive information, it is provided by multiple sensors, such as the individual joint encoders and force senors on each foot, meaning that they can also be damaged independently due to wear and tear from daily operations and it is not easy to have a redundant sensor. Among these sensors, joint encoders are the source of both q and q̇ for the observation. From previous analysis, missing of joint information can be crucial. We investigate the scenario where only a few encoders are dead or the data are compromised and need to be excluded. We conduct the test to mask certain numbers of joint encoders. For each masked joint, the related q and q̇ are removed completely from the observation. The results are shown in Figure <ref>. The loss of the joint encoder can have a great impact on the performance as the related information is very essential for quadruped locomotion. However, our transformer model can still handle multiple missing encoders before large performance degradation. 








Time Window. Another configuration we can easily adjust on the fly for balance is the time window for historical information. The default window, T=15, is equivalently to past 0.3s. We check whether such a long sequence of information is necessary by applying different time windows without masking. The results are shown in Figure <ref>. It is clear that  can efficiently extract and reconstruct the robot state for actions with only 7 steps of past information. Interestingly, given a longer timeframe like T=20 and the position the transformer has never seen during training,  is still robust and not affected by these unknown information.   








Unseen Information. We further assess the model's capability to handle previously unseen information by appending the sequence with random dummy tokens. With the increasing of unseen information, it become more challenging for the policy in following commands. However, even with 256 dummy tokens, equivalent to a camera frame with ViT <cit.>, the agent can still produce reasonable commands for exploration. This demonstrates that the model can be used as a solid foundation for further extension with high-dimensional information.

Minimized Observation and Fine-tuning. From the previous analysis in this section, we have identify the important sensors and the minimal history length required. We further explore the feasibility of creating a minimized observation policy based on these information. A observation with only cmd, q, g and a_t-1 with a window of T=7 is equivalent to a mask of 71% on the complete observation space. When directly deployed with such mask, the policy cannot perform well due to all the missing information. As there is no more observation, we freeze the transformer for a fast finetuning on the projection layers and test both the vanilla PPO <cit.> and supervised learning with online correction <cit.>. The performance of the policies is shown in Figure <ref>. While both algorithms can help improve the performance of the policy with only minimized observations, supervised learning gives a more significant boost. Training with the teacher has been identified as one major approach to achieve quadruped locomotion on challenging terrains <cit.>. Although our foundation with the transformer can provide a solid start point of student policy, additional work is still needed for pure RL-based fine-tuning to reduce the dependence on privilege information. 








Model Size. We adjust our model by stacking [1, 2, 3, 4] attention blocks, resulting in a model with 0.2M, 0.33M, 0.46M and 0.59M parameters respectively. Figure <ref> shows that the model performance increases with the model size reaching optimal results with three stacked layers.












 §.§ Physical Deployment

We successfully deployed the trained policy, exported with JIT, directly on a Unitree A1 robot equipped with a Jetson AGX Orin Developer Kit. The Jetson acts as both the main processor and a payload. No further model optimization is required for a zero-shot transfer. With the onboard processing power, the policy can process the full observation space at 75Hz and 100Hz while masking out half of the data. With our minimized observation, it can achieve 150Hz, meeting the requirements for real-time deployment and allowing room for further extension with high-dimensional sensors.  Figure <ref> shows some snapshots from the deployment test. Please refer to the supplementary video for more information.



§ CONCLUSION


This paper introduces , a transformer-based model for quadruped locomotion. It leverages the masking technique and the direct sensor-level attention to enhance the understanding and generation of sensory information input. We evaluate the robustness of our policy with different combinations of proprioceptive information and demonstrate its capability to compensate for missing data and handle unseen information. Finally, we show that our model is efficient to be deployed on a physical robot without any additional optimization. 

Limitation and Future Work. Attention in the full sensory-temporal observation space is computationally intensive and time-consuming. Although using masking can significantly reduce the resources needed, it still takes hours for knowledge transfer, which is considerably longer than existing methods like TCN and temporal-level attention. Additionally, fine-tuning the model with pure reinforcement learning remains challenging, necessitating a more efficient knowledge transfer solution to leverage privileged information effectively. While the policy demonstrated its capability to handle missing data, an addition module is needed to detect and mask out the defected sensors.







  





§ QUADRUPED SENSOR SET


When selecting sensors for proprioceptive information gathering, we referred to some of the most popular models that are commonly used. Table <ref> summarizes the default sensor sets that are typically included out of the box. For low-level tasks such as locomotion, joint encoders and IMUs are the most commonly used sensors, with foot contact sensors also being utilized in some works <cit.>.





§ DOMAIN RANDOMIZATION


We provide details of the domain randomization techniques used in simulation in Table <ref>. To simulate user commands, we frequently resample linear commands in longitudinal and lateral directions, as well as the angular heading, which are then converted to angular velocity commands. To enhance the robustness and generalization of our policy, key robot control factors such as stiffness, damping, and motor strength are randomized. External interferences, such as payload variations and push forces, are also randomized. To accommodate the processing time of the transformer, each action can be delayed by up to 15 ms. The same randomization techniques are applied to both the teacher training session and the knowledge transfer session. 









§ REWARD TERMS


Following the approaches in  <cit.>, the rewards are designed to encourage the agent to follow the commanded velocity. To achieve a smooth, safe, and efficient policy, we penalize unwanted motion, sudden changes in actions, feet slippage, body collisions, and excessive power consumption. During testing, we noticed that adding a reward component for balanced power distribution across joints helps to achieve more balanced motion.





§ TRAINING HYPER-PARAMETERS


In Table <ref>, present the detailed training hyper-parameters for PPO training of the teacher policy and the knowledge transfer for transformer-based student policies. The two stages in the knowledge transfer process share the same hyper-parameters. 





§ EXTENSION WITH HEIGHT INFORMATION


To better illustrate our framework's capabilities, we tokenize height map information using a vanilla MLP encoder and add it directly to the observation for fine-tuning over 200 epochs. The normalized test results are shown in Table <ref>. The inclusion of height information significantly improves the agent's ability to navigate through terrains, particularly on staircases, which pose the greatest challenge for a small quadruped.






§ LEARNABLE MASK


We conduct additional tests to compare the performance of using a learnable representation with the same two-stage transfer. The results, shown in Table <ref>, indicate that when deployed with masks, the learnable representation cannot match the performance of the removal approach and has a significantly negative impact when all information is present.



