35 lines
1.4 KiB
Python
35 lines
1.4 KiB
Python
from legged_gym.envs.go2.go2_env import Go2Robot
|
|
import torch
|
|
|
|
|
|
class Go1Robot(Go2Robot):
|
|
"""Go1 asset with an explicit policy-to-asset joint permutation.
|
|
|
|
Isaac Gym loads this asset as [FL, FR, RL, RR]. The permutation below
|
|
presents observations to the policy as [FR, FL, RR, RL] and maps policy
|
|
actions back to the Isaac Gym asset order. Existing Go1 checkpoints depend
|
|
on this policy order.
|
|
"""
|
|
|
|
# Policy order [FR, FL, RR, RL] -> asset order [FL, FR, RL, RR].
|
|
_POLICY_TO_ASSET = [3, 4, 5, 0, 1, 2, 9, 10, 11, 6, 7, 8]
|
|
|
|
def step(self, actions):
|
|
asset_actions = torch.zeros_like(actions)
|
|
asset_actions[:, self._POLICY_TO_ASSET] = actions
|
|
return super().step(asset_actions)
|
|
|
|
def compute_observations(self):
|
|
super().compute_observations()
|
|
|
|
# Go2 policy observation: angular velocity, gravity, commands,
|
|
# joint-position error, joint velocity, previous action.
|
|
for start in (9, 21, 33):
|
|
self.obs_buf[:, start:start + 12] = self.obs_buf[:, start:start + 12][:, self._POLICY_TO_ASSET]
|
|
|
|
# Keep the privileged teacher input consistent as well. Its first
|
|
# three fields are base linear velocity, so joint fields start at 12.
|
|
if self.privileged_obs_buf is not None:
|
|
for start in (12, 24, 36, 52, 64):
|
|
self.privileged_obs_buf[:, start:start + 12] = self.privileged_obs_buf[:, start:start + 12][:, self._POLICY_TO_ASSET]
|