bigym/bigym/envs/cupboards.py at master · NeuracoreAI/bigym · GitHub

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
"""Cupboard interaction tasks."""
from abc import ABC

import numpy as np

from bigym.bigym_env import BiGymEnv
from bigym.const import PRESETS_PATH
from bigym.envs.props.cabintets import BaseCabinet, WallCabinet


TOLERANCE = 0.1


class _CupboardsInteractionEnv(BiGymEnv, ABC):
    """Base cupboards environment."""

    RESET_ROBOT_POS = np.array([-0.2, 0, 0])

    _PRESET_PATH = PRESETS_PATH / "counter_base_wall_3x1.yaml"

    def _initialize_env(self):
        self.cabinet_drawers = self._preset.get_props(BaseCabinet)[0]
        self.cabinet_door_left = self._preset.get_props(BaseCabinet)[1]
        self.cabinet_door_right = self._preset.get_props(BaseCabinet)[2]
        self.cabinet_wall = self._preset.get_props(WallCabinet)[0]
        self.all_cabinets = [
            self.cabinet_drawers,
            self.cabinet_door_left,
            self.cabinet_door_right,
            self.cabinet_wall,
        ]


class DrawerTopOpen(_CupboardsInteractionEnv):
    """Open top drawer of the cupboard task."""

    def _success(self) -> bool:
        return np.allclose(self.cabinet_drawers.get_state()[-1], 1, atol=TOLERANCE)


class DrawerTopClose(_CupboardsInteractionEnv):
    """Close top drawer of the cupboard task."""

    def _success(self) -> bool:
        return np.allclose(self.cabinet_drawers.get_state()[-1], 0, atol=TOLERANCE)

    def _on_reset(self):
        self.cabinet_drawers.set_state(np.array([0, 0, 1]))


class DrawersAllOpen(_CupboardsInteractionEnv):
    """Open all drawers of the cupboard task."""

    def _success(self) -> bool:
        return np.allclose(self.cabinet_drawers.get_state(), 1, atol=TOLERANCE)


class DrawersAllClose(_CupboardsInteractionEnv):
    """Close all drawers of the cupboard task."""

    def _success(self) -> bool:
        return np.allclose(self.cabinet_drawers.get_state(), 0, atol=TOLERANCE)

    def _on_reset(self):
        self.cabinet_drawers.set_state(np.array([1, 1, 1]))


class WallCupboardOpen(_CupboardsInteractionEnv):
    """Open doors of the wall cupboard task."""

    def _success(self) -> bool:
        return np.allclose(self.cabinet_wall.get_state(), 1, atol=TOLERANCE)


class WallCupboardClose(_CupboardsInteractionEnv):
    """Close doors of the wall cupboard task."""

    def _success(self) -> bool:
        return np.allclose(self.cabinet_wall.get_state(), 0, atol=TOLERANCE)

    def _on_reset(self):
        self.cabinet_wall.set_state(np.array([1, 1]))


class CupboardsOpenAll(_CupboardsInteractionEnv):
    """Open all doors/drawers of the kitchen counter task."""

    def _success(self) -> bool:
        for cabinet in self.all_cabinets:
            if not np.allclose(cabinet.get_state(), 1, atol=TOLERANCE):
                return False
        return True


class CupboardsCloseAll(_CupboardsInteractionEnv):
    """Close all doors/drawers of the kitchen counter task."""

    def _success(self) -> bool:
        for cabinet in self.all_cabinets:
            if not np.allclose(cabinet.get_state(), 0, atol=TOLERANCE):
                return False
        return True

    def _on_reset(self):
        for cabinet in self.all_cabinets:
            cabinet.set_state(np.ones_like(cabinet.get_state()))