Spaces:
Runtime error
Runtime error
import numpy as np | |
from cliport.tasks.task import Task | |
from cliport.utils import utils | |
class TowersOfHanoi(Task): | |
"""Sequentially move disks from one tower to another—only smaller disks can be on top of larger ones.""" | |
def __init__(self): | |
super().__init__() | |
self.max_steps = 14 | |
self.lang_template = "solve towers of hanoi" | |
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') | |
# Rod positions in base coordinates. | |
rod_position = ((0, -0.12, 0.03), (0, 0, 0.03), (0, 0.12, 0.03)) | |
# Add disks. | |
disks = [] | |
n_disks = 3 | |
for i in range(n_disks): | |
disk_urdf = 'hanoi/disk%d.urdf' % i | |
pos = utils.apply(base_pose, rod_position[0]) | |
z = 0.015 * (n_disks - i - 2) | |
pos = (pos[0], pos[1], pos[2] + z) | |
disks.append(env.add_object(disk_urdf, (pos, base_pose[1]))) | |
# 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_position = rod_position[step[2]] | |
targ_position = utils.apply(base_pose, targ_position) | |
targ_pose = (targ_position, (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) |