|
14 | 14 | RCSpaceType, |
15 | 15 | Vec6Type, |
16 | 16 | Vec7Type, |
| 17 | + Vec18Type, |
17 | 18 | VecType, |
18 | 19 | get_space, |
19 | 20 | get_space_keys, |
20 | 21 | ) |
| 22 | +from rcs.hand.interface import BaseHand |
21 | 23 |
|
22 | 24 | _logger = logging.getLogger(__name__) |
23 | 25 |
|
@@ -99,6 +101,22 @@ class GripperDictType(RCSpaceType): |
99 | 101 | gripper: Annotated[float, gym.spaces.Box(low=0, high=1, dtype=np.float32)] |
100 | 102 |
|
101 | 103 |
|
| 104 | +class HandBinDictType(RCSpaceType): |
| 105 | + # 0 for closed, 1 for open (>=0.5 for open) |
| 106 | + hand: Annotated[float, gym.spaces.Box(low=0, high=1, dtype=np.float32)] |
| 107 | + |
| 108 | + |
| 109 | +class HandVecDictType(RCSpaceType): |
| 110 | + hand: Annotated[ |
| 111 | + Vec18Type, |
| 112 | + gym.spaces.Box( |
| 113 | + low=np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), |
| 114 | + high=np.array([1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]), |
| 115 | + dtype=np.float32, |
| 116 | + ), |
| 117 | + ] |
| 118 | + |
| 119 | + |
102 | 120 | class CameraDictType(RCSpaceType): |
103 | 121 | frames: dict[ |
104 | 122 | Annotated[str, "camera_names"], |
@@ -627,3 +645,74 @@ def action(self, action: dict[str, Any]) -> dict[str, Any]: |
627 | 645 | self._last_gripper_cmd = gripper_action |
628 | 646 | del action[self.gripper_key] |
629 | 647 | return action |
| 648 | + |
| 649 | + |
| 650 | +class HandWrapper(ActObsInfoWrapper): |
| 651 | + """ |
| 652 | + This wrapper allows for controlling the hand of the robot |
| 653 | + using either binary or continuous actions. |
| 654 | + The binary action space allows for opening and closing the hand, |
| 655 | + while the continuous action space allows for setting the hand |
| 656 | + to a specific pose. |
| 657 | + The wrapper also provides an observation space that includes |
| 658 | + the hand state. |
| 659 | + The hand state is represented as a binary value (0 for closed, |
| 660 | + 1 for open) or as a continuous value (normalized joint positions). |
| 661 | + """ |
| 662 | + |
| 663 | + BINARY_HAND_CLOSED = 0 |
| 664 | + BINARY_HAND_OPEN = 1 |
| 665 | + |
| 666 | + def __init__(self, env, hand: BaseHand, binary: bool = True): |
| 667 | + super().__init__(env) |
| 668 | + self.unwrapped: RobotEnv |
| 669 | + self.observation_space: gym.spaces.Dict |
| 670 | + self.action_space: gym.spaces.Dict |
| 671 | + self.binary = binary |
| 672 | + if self.binary: |
| 673 | + self.observation_space.spaces.update(get_space(HandBinDictType).spaces) |
| 674 | + self.action_space.spaces.update(get_space(HandBinDictType).spaces) |
| 675 | + self.hand_key = get_space_keys(HandBinDictType)[0] |
| 676 | + else: |
| 677 | + self.observation_space.spaces.update(get_space(HandVecDictType).spaces) |
| 678 | + self.action_space.spaces.update(get_space(HandVecDictType).spaces) |
| 679 | + self.hand_key = get_space_keys(HandVecDictType)[0] |
| 680 | + self._hand = hand |
| 681 | + self._last_hand_cmd = None |
| 682 | + |
| 683 | + def reset(self, **kwargs) -> tuple[dict[str, Any], dict[str, Any]]: |
| 684 | + self._hand.reset() |
| 685 | + self._last_hand_cmd = None |
| 686 | + return super().reset(**kwargs) |
| 687 | + |
| 688 | + def observation(self, observation: dict[str, Any], info: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]: |
| 689 | + observation = copy.deepcopy(observation) |
| 690 | + if self.binary: |
| 691 | + observation[self.hand_key] = ( |
| 692 | + self._last_hand_cmd if self._last_hand_cmd is not None else self.BINARY_HAND_OPEN |
| 693 | + ) |
| 694 | + else: |
| 695 | + observation[self.hand_key] = self._hand.get_normalized_joints_poses() |
| 696 | + |
| 697 | + info = {} |
| 698 | + return observation, info |
| 699 | + |
| 700 | + def action(self, action: dict[str, Any]) -> dict[str, Any]: |
| 701 | + |
| 702 | + action = copy.deepcopy(action) |
| 703 | + assert self.hand_key in action, "hand action not found." |
| 704 | + |
| 705 | + hand_action = np.round(action[self.hand_key]) if self.binary else action[self.hand_key] |
| 706 | + hand_action = np.clip(hand_action, 0.0, 1.0) |
| 707 | + |
| 708 | + if self.binary: |
| 709 | + if self._last_hand_cmd is None or self._last_hand_cmd != hand_action: |
| 710 | + if hand_action == self.BINARY_HAND_CLOSED: |
| 711 | + self._hand.grasp() |
| 712 | + else: |
| 713 | + self._hand.open() |
| 714 | + else: |
| 715 | + self._hand.set_normalized_joints_poses(hand_action) |
| 716 | + self._last_hand_cmd = hand_action |
| 717 | + del action[self.hand_key] |
| 718 | + return action |
0 commit comments