Spaces:
Runtime error
Runtime error
| import numpy as np | |
| from cliport.tasks.task import Task | |
| from cliport.utils import utils | |
| import pybullet as p | |
| import random | |
| class TowersOfHanoiSeq(Task): | |
| """Move the ring to the specified peg in the language instruction at each timestep""" | |
| def __init__(self): | |
| super().__init__() | |
| self.max_steps = 14 | |
| self.lang_template = "move the {obj} ring to the {loc}" | |
| self.task_completed_desc = "solved towers of hanoi." | |
| self.additional_reset() | |
| def reset(self, env): | |
| super().reset(env) | |
| # Add stand. | |
| base_size = (0.12, 0.36, 0.01) | |
| base_urdf = 'hanoi/stand.urdf' | |
| base_pose = self.get_random_pose(env, base_size) | |
| env.add_object(base_urdf, base_pose, 'fixed') | |
| # Choose three colors for three rings. | |
| colors, color_names = utils.get_colors(mode=self.mode, n_colors=3) | |
| # Rod positions in base coordinates. | |
| rod_pos = ((0, -0.12, 0.03), (0, 0, 0.03), (0, 0.12, 0.03)) | |
| rod_names = ('lighter brown side', 'middle of the stand', 'darker brown side') | |
| # Add disks. | |
| disks = [] | |
| disks_names = {} | |
| n_disks = 3 | |
| for i in range(n_disks): | |
| disk_urdf = 'hanoi/disk%d.urdf' % i | |
| pos = utils.apply(base_pose, rod_pos[0]) | |
| z = 0.015 * (n_disks - i - 2) | |
| pos = (pos[0], pos[1], pos[2] + z) | |
| ring_id = env.add_object(disk_urdf, (pos, base_pose[1]), color=colors[i]) | |
| disks.append(ring_id) | |
| disks_names[ring_id] = color_names[i] | |
| # Solve Hanoi sequence with dynamic programming. | |
| hanoi_steps = utils.solve_hanoi_all(n_disks) | |
| # Goal: pick and place disks using Hanoi sequence. | |
| for step in hanoi_steps: | |
| disk_id = disks[step[0]] | |
| targ_pos = rod_pos[step[2]] | |
| targ_pos = utils.apply(base_pose, targ_pos) | |
| targ_pose = (targ_pos, (0, 0, 0, 1)) | |
| self.add_goal(objs=[disk_id], matches=np.int32([[1]]), targ_poses=[targ_pose], replace=False, | |
| rotations=True, metric='pose', params=None, step_max_reward=1 / len(hanoi_steps), | |
| symmetries=[0]) | |
| self.lang_goals.append(self.lang_template.format(obj=disks_names[disk_id], | |
| loc=rod_names[step[2]])) |