Project_MultiagentCliport/cliport/tasks/block_insertion.py at main · AutonoBot-Lab/Project_MultiagentCliport · 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
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
"""Insertion Tasks."""

import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils

import pybullet as p


class BlockInsertion(Task):
    """Insertion Task - Base Variant."""

    def __init__(self):
        super().__init__()
        self.max_steps = 3
        self.lang_template = "put the L shape block in the L shape hole"
        self.task_completed_desc = "done with insertion."

    def reset(self, env):
        super().reset(env)
        block_id = self.add_block(env)
        targ_pose = self.add_fixture(env)
        # self.goals.append(
        #     ([block_id], [2 * np.pi], [[0]], [targ_pose], 'pose', None, 1.))
        self.goals.append(([(block_id, (2 * np.pi, None))], np.int32([[1]]),
                           [targ_pose], False, True, 'pose', None, 1))
        self.lang_goals.append(self.lang_template)

    def add_block(self, env):
        """Add L-shaped block."""
        size = (0.1, 0.1, 0.04)
        urdf = 'insertion/ell.urdf'
        pose = self.get_random_pose(env, size)
        return env.add_object(urdf, pose)

    def add_fixture(self, env):
        """Add L-shaped fixture to place block."""
        size = (0.1, 0.1, 0.04)
        urdf = 'insertion/fixture.urdf'
        pose = self.get_random_pose(env, size)
        env.add_object(urdf, pose, 'fixed')
        return pose


class BlockInsertionTranslation(BlockInsertion):
    """Insertion Task - Translation Variant."""

    def get_random_pose(self, env, obj_size):
        pose = super(BlockInsertionTranslation, self).get_random_pose(env, obj_size)
        pos, rot = pose
        rot = utils.eulerXYZ_to_quatXYZW((0, 0, np.pi / 2))
        return pos, rot

    # Visualization positions.
    # block_pos = (0.40, -0.15, 0.02)
    # fixture_pos = (0.65, 0.10, 0.02)


class BlockInsertionEasy(BlockInsertionTranslation):
    """Insertion Task - Easy Variant."""

    def add_block(self, env):
        """Add L-shaped block in fixed position."""
        # size = (0.1, 0.1, 0.04)
        urdf = 'insertion/ell.urdf'
        pose = ((0.5, 0, 0.02), p.getQuaternionFromEuler((0, 0, np.pi / 2)))
        return env.add_object(urdf, pose)


class BlockInsertionSixDof(BlockInsertion):
    """Insertion Task - 6DOF Variant."""

    def __init__(self):
        super().__init__()
        self.sixdof = True
        self.pos_eps = 0.02

    def add_fixture(self, env):
        """Add L-shaped fixture to place block."""
        size = (0.1, 0.1, 0.04)
        urdf = 'insertion/fixture.urdf'
        pose = self.get_random_pose_6dof(env, size)
        env.add_object(urdf, pose, 'fixed')
        return pose

    def get_random_pose_6dof(self, env, obj_size):
        pos, rot = super(BlockInsertionSixDof, self).get_random_pose(env, obj_size)
        z = (np.random.rand() / 10) + 0.03
        pos = (pos[0], pos[1], obj_size[2] / 2 + z)
        roll = (np.random.rand() - 0.5) * np.pi / 2
        pitch = (np.random.rand() - 0.5) * np.pi / 2
        yaw = np.random.rand() * 2 * np.pi
        rot = utils.eulerXYZ_to_quatXYZW((roll, pitch, yaw))
        return pos, rot


class BlockInsertionNoFixture(BlockInsertion):
    """Insertion Task - No Fixture Variant."""

    def add_fixture(self, env):
        """Add target pose to place block."""
        size = (0.1, 0.1, 0.04)
        # urdf = 'insertion/fixture.urdf'
        pose = self.get_random_pose(env, size)
        return pose

    # def reset(self, env, last_info=None):
    #   self.num_steps = 1
    #   self.goal = {'places': {}, 'steps': []}

    #   # Add L-shaped block.
    #   block_size = (0.1, 0.1, 0.04)
    #   block_urdf = 'insertion/ell.urdf'
    #   block_pose = self.get_random_pose(env, block_size)
    #   block_id = env.add_object(block_urdf, block_pose)
    #   self.goal['steps'].append({block_id: (2 * np.pi, [0])})

    #   # Add L-shaped target pose, but without actually adding it.
    #   if self.goal_cond_testing:
    #     assert last_info is not None
    #     self.goal['places'][0] = self._get_goal_info(last_info)
    #     # print('\nin insertion reset, goal: {}'.format(self.goal['places'][0]))
    #   else:
    #     hole_pose = self.get_random_pose(env, block_size)
    #     self.goal['places'][0] = hole_pose
    #     # print('\nin insertion reset, goal: {}'.format(hole_pose))

    # def _get_goal_info(self, last_info):
    #   """Used to determine the goal given the last `info` dict."""
    #   position, rotation, _ = last_info[4]  # block ID=4
    #   return (position, rotation)