-
Notifications
You must be signed in to change notification settings - Fork 19
/
Copy pathdishwasher_cups.py
131 lines (106 loc) · 4.54 KB
/
dishwasher_cups.py
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
"""Load/unload cups to/from dishwasher."""
from abc import ABC
import numpy as np
from pyquaternion import Quaternion
from bigym.bigym_env import BiGymEnv
from bigym.const import PRESETS_PATH
from bigym.envs.props.cabintets import BaseCabinet, WallCabinet
from bigym.envs.props.dishwasher import Dishwasher
from bigym.envs.props.kitchenware import Mug
from bigym.utils.env_utils import get_random_sites
class _DishwasherCupsEnv(BiGymEnv, ABC):
"""Base cups environment."""
RESET_ROBOT_POS = np.array([0, -0.6, 0])
_PRESET_PATH = PRESETS_PATH / "counter_dishwasher.yaml"
_CUPS_COUNT = 2
def _initialize_env(self):
self.dishwasher = self._preset.get_props(Dishwasher)[0]
self.cabinets = self._preset.get_props(BaseCabinet)
self.cups = [Mug(self._mojo) for _ in range(self._CUPS_COUNT)]
def _fail(self) -> bool:
if super()._fail():
return True
for cup in self.cups:
if cup.is_colliding(self.floor):
return True
return False
def _on_reset(self):
self.dishwasher.set_state(door=1, bottom_tray=0, middle_tray=1)
class DishwasherUnloadCups(_DishwasherCupsEnv):
"""Unload cups from dishwasher task."""
_SITES_STEP = 3
_SITES_SLICE = 3
_CUPS_ROT_X = np.deg2rad(180)
_CUPS_ROT_Z = np.deg2rad(90)
_CUPS_ROT_BOUNDS = np.deg2rad(5)
_CUPS_POS = np.array([0, -0.05, 0.05])
_CUPS_STEP = np.array([0.115, 0, 0])
def _success(self) -> bool:
for cup in self.cups:
if not any([cup.is_colliding(cabinet) for cabinet in self.cabinets]):
return False
for side in self.robot.grippers:
if self.robot.is_gripper_holding_object(cup, side):
return False
return True
def _on_reset(self):
super()._on_reset()
sites = self.dishwasher.tray_middle.site_sets[0]
sites = get_random_sites(
sites, len(self.cups), self._SITES_STEP, self._SITES_SLICE
)
for site, cup in zip(sites, self.cups):
quat = Quaternion(axis=[1, 0, 0], angle=self._CUPS_ROT_X)
angle = np.random.uniform(-self._CUPS_ROT_BOUNDS, self._CUPS_ROT_BOUNDS)
quat *= Quaternion(axis=[0, 0, 1], angle=self._CUPS_ROT_Z + angle)
cup.body.set_quaternion(quat.elements, True)
pos = site.get_position()
pos += self._CUPS_POS
cup.body.set_position(pos, True)
class DishwasherUnloadCupsLong(DishwasherUnloadCups):
"""Unload cup from dishwasher in wall cabinet task."""
_PRESET_PATH = PRESETS_PATH / "counter_dishwasher_wall_cabinet.yaml"
_CUPS_COUNT = 1
_SITES_SLICE = 2
_TOLERANCE = 0.1
def _initialize_env(self):
super()._initialize_env()
self.wall_cabinet = self._preset.get_props(WallCabinet)[0]
def _success(self) -> bool:
if not np.allclose(self.dishwasher.get_state(), 0, atol=self._TOLERANCE):
return False
if not np.allclose(self.wall_cabinet.get_state(), 0, atol=self._TOLERANCE):
return False
for cup in self.cups:
if not cup.is_colliding(self.wall_cabinet.shelf_bottom):
return False
for side in self.robot.grippers:
if self.robot.is_gripper_holding_object(cup, side):
return False
return True
class DishwasherLoadCups(_DishwasherCupsEnv):
"""Load cups to dishwasher task."""
_CUPS_POS = np.array([0.6, -0.6, 1])
_CUPS_POS_STEP = np.array([0, 0.15, 0])
_CUPS_POS_BOUNDS = np.array([0.05, 0.02, 0])
_CUPS_ROT_X = np.deg2rad(180)
_CUPS_ROT_Z = np.deg2rad(180)
_CUPS_ROT_BOUNDS = np.deg2rad(30)
def _success(self) -> bool:
for cup in self.cups:
if not cup.is_colliding(self.dishwasher.tray_middle.colliders):
return False
for side in self.robot.grippers:
if self.robot.is_gripper_holding_object(cup, side):
return False
return True
def _on_reset(self):
super()._on_reset()
for i, cup in enumerate(self.cups):
quat = Quaternion(axis=[1, 0, 0], angle=self._CUPS_ROT_X)
angle = np.random.uniform(-self._CUPS_ROT_BOUNDS, self._CUPS_ROT_BOUNDS)
quat *= Quaternion(axis=[0, 0, 1], angle=self._CUPS_ROT_Z + angle)
cup.body.set_quaternion(quat.elements, True)
pos = self._CUPS_POS + i * self._CUPS_POS_STEP
pos += np.random.uniform(-self._CUPS_POS_BOUNDS, self._CUPS_POS_BOUNDS)
cup.body.set_position(pos, True)