Spaces:
Runtime error
Runtime error
File size: 2,911 Bytes
8fc2b4e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
import numpy as np
import os
import pybullet as p
import random
from cliport.tasks import primitives
from cliport.tasks.grippers import Spatula
from cliport.tasks.task import Task
from cliport.utils import utils
import numpy as np
from cliport.tasks.task import Task
from cliport.utils import utils
class AssembleSingleCar(Task):
"""Assemble a mini car using a large blue box as the body, a smaller red box on top as the roof, and two tiny green boxes on the sides as wheels."""
def __init__(self):
super().__init__()
self.max_steps = 10
self.lang_template = "build a mini car using a large blue box as the body, a smaller red box on top as the roof, and two tiny green boxes on the sides as wheels"
self.task_completed_desc = "done assembling the car."
def reset(self, env):
super().reset(env)
# Add car body (large blue box).
body_size = (0.1, 0.05, 0.02) # x, y, z dimensions
body_pose = self.get_random_pose(env, body_size)
body_urdf = 'box/box-template.urdf'
body_color = utils.COLORS['blue']
body_id = env.add_object(body_urdf, body_pose, color=body_color)
# Add car roof (smaller red box).
roof_size = (0.08, 0.04, 0.02) # x, y, z dimensions
roof_pose = self.get_random_pose(env, roof_size)
roof_urdf = 'box/box-template.urdf'
roof_color = utils.COLORS['red']
roof_id = env.add_object(roof_urdf, roof_pose, color=roof_color)
# Add car wheels (two tiny green boxes).
wheel_size = (0.02, 0.02, 0.01) # x, y, z dimensions
wheel_urdf = 'box/box-template.urdf'
wheel_color = utils.COLORS['green']
wheel_ids = []
for _ in range(2):
wheel_pose = self.get_random_pose(env, wheel_size)
wheel_id = env.add_object(wheel_urdf, wheel_pose, color=wheel_color)
wheel_ids.append(wheel_id)
# Goal: assemble the car by placing the roof on the body and the wheels on the sides.
# The target poses are calculated based on the body pose.
roof_targ_pose = (body_pose[0] + np.array([0, 0, body_size[2] + roof_size[2]/2]), body_pose[1])
wheel_targ_poses = [(body_pose[0] + np.array([0, body_size[1]/2 + wheel_size[1]/2, -body_size[2]/2]), body_pose[1]),
(body_pose[0] + np.array([0, -body_size[1]/2 - wheel_size[1]/2, -body_size[2]/2]), body_pose[1])]
# Add the goals.
self.add_goal(objs=[roof_id], matches=np.ones((1, 1)), targ_poses=[roof_targ_pose], replace=False,
rotations=True, metric='pose', params=None, step_max_reward=1/3)
self.add_goal(objs=wheel_ids, matches=np.ones((2, 2)), targ_poses=wheel_targ_poses, replace=False,
rotations=True, metric='pose', params=None, step_max_reward=2/3)
self.lang_goals.append(self.lang_template) |