metadata
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:
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()
)