GenSim / cliport /tasks /generated_task.py
LeroyWaa's picture
add gensim code
8fc2b4e
import numpy as np
from cliport.tasks import Task
from cliport.tasks.grippers import Spatula
import random
from cliport.utils import utils
import pybullet as p
class GeneratedTask(Task):
"""Build a car using blocks."""
def __init__(self):
super().__init__()
self.max_steps = 20
self.lang_template = "build a car using blocks"
self.task_completed_desc = "done building car."
self.additional_reset()
def reset(self, env):
super().reset(env)
# Add car base.
base_size = (0.12, 0.12, 0.02)
base_urdf = 'car-base/car-base.urdf'
base_pose = self.get_random_pose(env, base_size)
env.add_object(base_urdf, base_pose, category='fixed')
# Block colors.
colors = [
utils.COLORS['red'], utils.COLORS['blue'], utils.COLORS['green'], utils.COLORS['yellow']
]
# Add blocks.
block_size = (0.04, 0.04, 0.04)
block_urdf = 'block/block.urdf'
objs = []
for i in range(4):
block_pose = self.get_random_pose(env, block_size)
block_id = env.add_object(block_urdf, block_pose, color=colors[i])
objs.append(block_id)
# Associate placement locations for goals.
place_pos = [(0, -0.05, 0.03), (0, 0, 0.03),
(0, 0.05, 0.03), (0, -0.025, 0.08)]
targs = [(utils.apply(base_pose, i), base_pose[1]) for i in place_pos]
# Goal: blocks are stacked vertically (bottom row: red, blue, green).
for i in range(4):
self.add_goal(objs=[objs[i]], matches=np.ones((1, 1)), targ_poses=[targs[i]], replace=False,
rotations=True, metric='pose', params=None, step_max_reward=1 / 4, symmetries=[np.pi/2])
self.lang_goals.append(self.lang_template.format(blocks="the blocks", stacked="stacked vertically"))
# Goal: blocks are stacked horizontally (top row: yellow).
for i in range(2):
self.add_goal(objs=[objs[i]], matches=np.ones((1, 1)), targ_poses=[targs[i]], replace=False,
rotations=True, metric='pose', params=None, step_max_reward=1 / 4)