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