GenSim / cliport /tasks /towers_of_hanoi.py
LeroyWaa's picture
add gensim code
8fc2b4e
raw
history blame
1.91 kB
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)