---
tags:
- model_hub_mixin
- pytorch_model_hub_mixin
pipeline_tag: reinforcement-learning
---

# This is an S5 based RL locomotion controller for the Unitree GO1 robot. Given checkpoint achieves:
- Linear velocity command from -2.0 m/sec to 2.0 m/sec in the forward direction (X).
- Linear velocity command from -1.2 m/sec to 1.2 m/sec in the side direction (Y).
- Angular velocity command from -1.2 rad/sec to 1.2 rad/sec around the body yaw direction (Rz).
- Height command from 0.10 m to 0.30 m in the up direction (Z).

# Model structure:
```python
ActorS5(
  (memory_proj): Linear(in_features=57, out_features=128, bias=True)
  (memory): S5Block(
    (s5): S5FAST(
      (seq): S5SSM()
    )
    (attn_norm): LayerNorm((128,), eps=1e-05, elementwise_affine=True)
    (geglu): GEGLU()
    (ff_enc): Linear(in_features=128, out_features=512, bias=False)
    (ff_dec): Linear(in_features=256, out_features=128, bias=False)
    (ff_norm): LayerNorm((128,), eps=1e-05, elementwise_affine=True)
  )
  (action_head): ActorGaussian(
    (actor): Sequential(
      (0): Linear(in_features=128, out_features=256, bias=True)
      (1): ELU(alpha=1.0)
      (2): Linear(in_features=256, out_features=128, bias=True)
      (3): ELU(alpha=1.0)
      (4): Linear(in_features=128, out_features=12, bias=True)
    )
    (_normalizer): Identity()
  )
  (_normalizer): EmpiricalNormalization()
)
```