You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
67 lines
2.6 KiB
67 lines
2.6 KiB
from typing import Any, Dict
|
|
|
|
import gymnasium as gym
|
|
import numpy as np
|
|
|
|
from decoupled_wbc.control.base.env import Env
|
|
from decoupled_wbc.control.envs.g1.utils.command_sender import BodyCommandSender
|
|
from decoupled_wbc.control.envs.g1.utils.state_processor import BodyStateProcessor
|
|
|
|
|
|
class G1Body(Env):
|
|
def __init__(self, config: Dict[str, Any]):
|
|
super().__init__()
|
|
self.body_state_processor = BodyStateProcessor(config=config)
|
|
self.body_command_sender = BodyCommandSender(config=config)
|
|
|
|
def observe(self) -> dict[str, any]:
|
|
body_state = self.body_state_processor._prepare_low_state() # (1, 148)
|
|
assert body_state.shape == (1, 148)
|
|
body_q = body_state[
|
|
0, 7 : 7 + 12 + 3 + 7 + 7
|
|
] # leg (12) + waist (3) + left arm (7) + right arm (7)
|
|
body_dq = body_state[0, 42 : 42 + 12 + 3 + 7 + 7]
|
|
body_ddq = body_state[0, 112 : 112 + 12 + 3 + 7 + 7]
|
|
body_tau_est = body_state[0, 77 : 77 + 12 + 3 + 7 + 7]
|
|
floating_base_pose = body_state[0, 0:7]
|
|
floating_base_vel = body_state[0, 36:42]
|
|
floating_base_acc = body_state[0, 106:112]
|
|
torso_quat = body_state[0, 141:145]
|
|
torso_ang_vel = body_state[0, 145:148]
|
|
|
|
return {
|
|
"body_q": body_q,
|
|
"body_dq": body_dq,
|
|
"body_ddq": body_ddq,
|
|
"body_tau_est": body_tau_est,
|
|
"floating_base_pose": floating_base_pose,
|
|
"floating_base_vel": floating_base_vel,
|
|
"floating_base_acc": floating_base_acc,
|
|
"torso_quat": torso_quat,
|
|
"torso_ang_vel": torso_ang_vel,
|
|
}
|
|
|
|
def queue_action(self, action: dict[str, any]):
|
|
# action should contain body_q, body_dq, body_tau
|
|
self.body_command_sender.send_command(
|
|
action["body_q"], action["body_dq"], action["body_tau"]
|
|
)
|
|
|
|
def observation_space(self) -> gym.Space:
|
|
return gym.spaces.Dict(
|
|
{
|
|
"body_q": gym.spaces.Box(low=-np.inf, high=np.inf, shape=(29,)),
|
|
"body_dq": gym.spaces.Box(low=-np.inf, high=np.inf, shape=(29,)),
|
|
"floating_base_pose": gym.spaces.Box(low=-np.inf, high=np.inf, shape=(7,)),
|
|
"floating_base_vel": gym.spaces.Box(low=-np.inf, high=np.inf, shape=(6,)),
|
|
}
|
|
)
|
|
|
|
def action_space(self) -> gym.Space:
|
|
return gym.spaces.Dict(
|
|
{
|
|
"body_q": gym.spaces.Box(low=-np.inf, high=np.inf, shape=(29,)),
|
|
"body_dq": gym.spaces.Box(low=-np.inf, high=np.inf, shape=(29,)),
|
|
"body_tau": gym.spaces.Box(low=-np.inf, high=np.inf, shape=(29,)),
|
|
}
|
|
)
|