| """ManiSkill-style peg insertion: put the peg into a square socket.""" |
| from ..registry import register_task |
| from ..envs import YamTaskEnv |
| from ..solvers import insert |
| from .. import conditions as C |
|
|
|
|
| @register_task("peg_insert") |
| class PegInsertTask(YamTaskEnv): |
| title = "PEG -> HOLE (insertion)" |
| tags = ["insertion", "precision", "primitive"] |
| instruction = {"default": "Put the peg in the hole.", |
| "specific": "Pick up the yellow peg and insert it into the square socket."} |
|
|
| peg_spawn = (-0.03, 0.10) |
| peg_spawn_jitter = (0.02, 0.015) |
| socket_xy = (-0.02, -0.22) |
| socket_jitter = (0.015, 0.015) |
| hole_width = 0.042 |
| |
| grasp_top = 0.03 |
| gripper_effort, gripper_damping = 75, 85 |
|
|
| def _load_scene(self): |
| self.scene.build_socket({"name": "socket", "xy": self.socket_xy, |
| "xy_jitter": self.socket_jitter, "hole": self.hole_width}) |
| self._placed = [self.scene.place_object( |
| {"name": "peg", "xy": self.peg_spawn, "xy_jitter": self.peg_spawn_jitter})] |
|
|
| def solve(self): |
| return insert.solve(self, obj="peg", target="socket", grasp_top=self.grasp_top) |
|
|
| def evaluate(self): |
| return self.check(C.labelled( |
| "peg seated in hole", C.object_seated_in_hole("peg", "socket", xy_tol=0.02, half_h=0.065))) |
|
|