File size: 8,759 Bytes
8cf4ead | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 | from typing import Any, Dict, SupportsFloat, Tuple
import gymnasium as gym
from gymnasium import spaces
import numpy as np
from decoupled_wbc.control.main.constants import DEFAULT_BASE_HEIGHT, DEFAULT_NAV_CMD
from decoupled_wbc.control.main.teleop.configs.configs import SyncSimDataCollectionConfig
from decoupled_wbc.control.policy.wbc_policy_factory import get_wbc_policy
from decoupled_wbc.control.robot_model import RobotModel
from decoupled_wbc.control.robot_model.instantiation import get_robot_type_and_model
class WholeBodyControlWrapper(gym.Wrapper):
"""Gymnasium wrapper to integrate whole-body control for locomotion/manipulation sims."""
def __init__(self, env, script_config):
super().__init__(env)
self.script_config = script_config
self.script_config["robot"] = env.unwrapped.robot_name
self.wbc_policy = self.setup_wbc_policy()
self._action_space = self._wbc_action_space()
@property
def robot_model(self) -> RobotModel:
"""Return the robot model from the wrapped environment."""
return self.env.unwrapped.robot_model # type: ignore
def reset(self, **kwargs):
obs, info = self.env.reset(**kwargs)
self.wbc_policy = self.setup_wbc_policy()
self.wbc_policy.set_observation(obs)
return obs, info
def step(self, action: Dict[str, Any]) -> Tuple[Any, SupportsFloat, bool, bool, Dict[str, Any]]:
action_dict = concat_action(self.robot_model, action)
wbc_goal = {}
for key in ["navigate_cmd", "base_height_command", "target_upper_body_pose"]:
if key in action_dict:
wbc_goal[key] = action_dict[key]
self.wbc_policy.set_goal(wbc_goal)
wbc_action = self.wbc_policy.get_action()
result = super().step(wbc_action)
self.wbc_policy.set_observation(result[0])
return result
def setup_wbc_policy(self):
robot_type, robot_model = get_robot_type_and_model(
self.script_config["robot"],
enable_waist_ik=self.script_config.get("enable_waist", False),
)
config = SyncSimDataCollectionConfig.from_dict(self.script_config)
config.update(
{
"save_img_obs": False,
"ik_indicator": False,
"enable_real_device": False,
"replay_data_path": None,
}
)
wbc_config = config.load_wbc_yaml()
wbc_config["upper_body_policy_type"] = "identity"
wbc_policy = get_wbc_policy(robot_type, robot_model, wbc_config, init_time=0.0)
self.total_dofs = len(robot_model.get_joint_group_indices("upper_body"))
wbc_policy.activate_policy()
return wbc_policy
def _get_joint_group_size(self, group_name: str) -> int:
"""Return the number of joints in a group, cached since lookup is static."""
if not hasattr(self, "_joint_group_size_cache"):
self._joint_group_size_cache = {}
if group_name not in self._joint_group_size_cache:
self._joint_group_size_cache[group_name] = len(
self.robot_model.get_joint_group_indices(group_name)
)
return self._joint_group_size_cache[group_name]
def _wbc_action_space(self) -> spaces.Dict:
action_space: Dict[str, spaces.Space] = {
"action.navigate_command": spaces.Box(
low=-np.inf, high=np.inf, shape=(3,), dtype=np.float32
),
"action.base_height_command": spaces.Box(
low=-np.inf, high=np.inf, shape=(1,), dtype=np.float32
),
"action.left_hand": spaces.Box(
low=-np.inf,
high=np.inf,
shape=(self._get_joint_group_size("left_hand"),),
dtype=np.float32,
),
"action.right_hand": spaces.Box(
low=-np.inf,
high=np.inf,
shape=(self._get_joint_group_size("right_hand"),),
dtype=np.float32,
),
"action.left_arm": spaces.Box(
low=-np.inf,
high=np.inf,
shape=(self._get_joint_group_size("left_arm"),),
dtype=np.float32,
),
"action.right_arm": spaces.Box(
low=-np.inf,
high=np.inf,
shape=(self._get_joint_group_size("right_arm"),),
dtype=np.float32,
),
}
if (
"waist"
in self.robot_model.supplemental_info.joint_groups["upper_body_no_hands"]["groups"] # type: ignore[attr-defined]
):
action_space["action.waist"] = spaces.Box(
low=-np.inf,
high=np.inf,
shape=(self._get_joint_group_size("waist"),),
dtype=np.float32,
)
return spaces.Dict(action_space)
def concat_action(robot_model: RobotModel, goal: Dict[str, Any]) -> Dict[str, Any]:
"""Combine individual joint-group targets into the upper-body action vector."""
processed_goal = {}
for key, value in goal.items():
processed_goal[key.replace("action.", "")] = value
first_value = next(iter(processed_goal.values()))
action = np.zeros(first_value.shape[:-1] + (robot_model.num_dofs,))
action_dict = {}
action_dict["navigate_cmd"] = processed_goal.pop("navigate_command", DEFAULT_NAV_CMD)
action_dict["base_height_command"] = np.array(
processed_goal.pop("base_height_command", DEFAULT_BASE_HEIGHT)
)
for joint_group, value in processed_goal.items():
indices = robot_model.get_joint_group_indices(joint_group)
action[..., indices] = value
upper_body_indices = robot_model.get_joint_group_indices("upper_body")
action = action[..., upper_body_indices]
action_dict["target_upper_body_pose"] = action
return action_dict
def prepare_observation_for_eval(robot_model: RobotModel, obs: dict) -> dict:
"""Add joint-group slices to an observation dict (real + sim evaluation helper)."""
assert "q" in obs, "q is not in the observation"
whole_q = obs["q"]
assert whole_q.shape[-1] == robot_model.num_joints, "q has wrong shape"
left_arm_q = whole_q[..., robot_model.get_joint_group_indices("left_arm")]
right_arm_q = whole_q[..., robot_model.get_joint_group_indices("right_arm")]
waist_q = whole_q[..., robot_model.get_joint_group_indices("waist")]
left_leg_q = whole_q[..., robot_model.get_joint_group_indices("left_leg")]
right_leg_q = whole_q[..., robot_model.get_joint_group_indices("right_leg")]
left_hand_q = whole_q[..., robot_model.get_joint_group_indices("left_hand")]
right_hand_q = whole_q[..., robot_model.get_joint_group_indices("right_hand")]
obs["state.left_arm"] = left_arm_q
obs["state.right_arm"] = right_arm_q
obs["state.waist"] = waist_q
obs["state.left_leg"] = left_leg_q
obs["state.right_leg"] = right_leg_q
obs["state.left_hand"] = left_hand_q
obs["state.right_hand"] = right_hand_q
return obs
def prepare_gym_space_for_eval(
robot_model: RobotModel, gym_space: gym.spaces.Dict
) -> gym.spaces.Dict:
"""Extend a gym Dict space with the joint-group keys used during evaluation."""
left_arm_space = spaces.Box(
low=-np.inf,
high=np.inf,
shape=(len(robot_model.get_joint_group_indices("left_arm")),),
)
right_arm_space = spaces.Box(
low=-np.inf,
high=np.inf,
shape=(len(robot_model.get_joint_group_indices("right_arm")),),
)
waist_space = spaces.Box(
low=-np.inf,
high=np.inf,
shape=(len(robot_model.get_joint_group_indices("waist")),),
)
left_leg_space = spaces.Box(
low=-np.inf,
high=np.inf,
shape=(len(robot_model.get_joint_group_indices("left_leg")),),
)
right_leg_space = spaces.Box(
low=-np.inf,
high=np.inf,
shape=(len(robot_model.get_joint_group_indices("right_leg")),),
)
left_hand_space = spaces.Box(
low=-np.inf,
high=np.inf,
shape=(len(robot_model.get_joint_group_indices("left_hand")),),
)
right_hand_space = spaces.Box(
low=-np.inf,
high=np.inf,
shape=(len(robot_model.get_joint_group_indices("right_hand")),),
)
gym_space["state.left_arm"] = left_arm_space
gym_space["state.right_arm"] = right_arm_space
gym_space["state.waist"] = waist_space
gym_space["state.left_leg"] = left_leg_space
gym_space["state.right_leg"] = right_leg_space
gym_space["state.left_hand"] = left_hand_space
gym_space["state.right_hand"] = right_hand_space
return gym_space
|