changed view of maniskill (d4de35c4) · Commits · 张泽凯 / VAGEN

vagen/env/primitive_skill/maniskill/env/align_two_cube.py

+11 −11

Original line number	Diff line number	Diff line
		@@ -22,6 +22,7 @@ class AlignTwoCubeEnv(BaseEnv):
		agent: Union[Panda, Xmate3Robotiq, Fetch]
		skill_config=None
		vlm_info_keys=[]
		state_keys=["red_cube_position", "green_cube_position"]

		def __init__(self, stage=0,args, robot_uids="panda", robot_init_qpos_noise=0.02, *kwargs):
		self.stage=stage
		@@ -29,7 +30,8 @@ class AlignTwoCubeEnv(BaseEnv):
		self.workspace_y=[-0.2, 0.2]
		self.workspace_z=[0.01, 0.2]
		self.robot_init_qpos_noise = robot_init_qpos_noise

		self.region_1 = np.array([[0.1, -0.2], [0.15, 0.1]])
		self.region_2 = np.array([[0.1, 0.1], [0.15, 0.2]])
		super().__init__(args, robot_uids=robot_uids, *kwargs)

		@property
		@@ -45,12 +47,12 @@ class AlignTwoCubeEnv(BaseEnv):

		@property
		def _default_sensor_configs(self):
		pose = sapien_utils.look_at(eye=[0.3, 0, 0.6], target=[-0.1, 0, 0.1])
		pose = sapien_utils.look_at(eye=[1, 0.0, 0.6], target=[-0.2, 0.0, 0.2])
		return [CameraConfig("base_camera", pose, 300, 300, np.pi / 2, 0.01, 100)]

		@property
		def _default_human_render_camera_configs(self):
		pose = sapien_utils.look_at([0.6, 0.7, 0.6], [0.0, 0.0, 0.35])
		pose = sapien_utils.look_at([1, 0.0, 0.6], [-0.2, 0.0, 0.2])
		return CameraConfig("render_camera", pose, 300,300, 1, 0.01, 100)

		def _load_scene(self, options: dict):
		@@ -74,10 +76,8 @@ class AlignTwoCubeEnv(BaseEnv):

		xyz = torch.zeros((b, 3))
		xyz[:, 2] = 0.02
		region_1 =[[self.workspace_x[0],self.workspace_y[0]],[self.workspace_x[0]+0.05,self.workspace_y[1]]]
		sampler_1 = randomization.UniformPlacementSampler(bounds=region_1, batch_size=b)
		region_2 =[[0.05,self.workspace_y[0]],[self.workspace_x[1],self.workspace_y[1]]]
		sampler_2 = randomization.UniformPlacementSampler(bounds=region_2, batch_size=b)
		sampler_1 = randomization.UniformPlacementSampler(bounds=self.region_1, batch_size=b)
		sampler_2 = randomization.UniformPlacementSampler(bounds=self.region_2, batch_size=b)
		radius = torch.linalg.norm(torch.tensor([0.02, 0.02])) + 0.02

		red_cube_xy = sampler_1.sample(radius, 100)
		@@ -127,8 +127,8 @@ class AlignTwoCubeEnv(BaseEnv):
		"is_red_cube_at_x0": is_red_cube_at_x0,
		"is_green_cube_grasped": is_green_cube_grasped,
		"is_green_cube_at_x0": is_green_cube_at_x0,
		"red_cube_pos": pos_A,
		"green_cube_pos": pos_B,
		"red_cube_position": pos_A,
		"green_cube_position": pos_B,
		"stage_0_success": is_red_cube_at_x0,
		"stage_1_success": is_green_cube_at_x0,
		"success": success.bool(),
		@@ -139,8 +139,8 @@ class AlignTwoCubeEnv(BaseEnv):
		def _get_obs_extra(self, info: Dict):
		assert "state" in self.obs_mode
		obs = dict(
		red_cube_pos=info["red_cube_pos"],
		green_cube_pos=info["green_cube_pos"],
		red_cube_position=info["red_cube_position"],
		green_cube_position=info["green_cube_position"],
		is_red_cube_at_x0=info["is_red_cube_at_x0"],
		is_green_cube_at_x0=info["is_green_cube_at_x0"],
		)

vagen/env/primitive_skill/maniskill/env/place_two_cube.py

+50 −48

Original line number	Diff line number	Diff line
		@@ -23,12 +23,14 @@ class PlaceTwoCubeEnv(BaseEnv):
		agent: Union[Panda, Xmate3Robotiq, Fetch]
		skill_config=None
		vlm_info_keys=[]
		state_keys=["red_cube_position", "green_cube_position"]

		def __init__(self, stage=0,args, robot_uids="panda", robot_init_qpos_noise=0.02, *kwargs):
		self.stage=stage
		self.workspace_x=[-0.10, 0.15]
		self.workspace_y=[-0.2, 0.2]
		self.workspace_z=[0.01, 0.2]
		self.region = np.array([[-0.1, -0.1], [0, 0.1]])
		self.robot_init_qpos_noise = robot_init_qpos_noise

		super().__init__(args, robot_uids=robot_uids, *kwargs)
		@@ -42,16 +44,16 @@ class PlaceTwoCubeEnv(BaseEnv):
		}

		def instruction(self):
		return "Please place cubeA at goalA and cubeB at goalB."
		return "Please place red cube at left target and green cube at right target."

		@property
		def _default_sensor_configs(self):
		pose = sapien_utils.look_at(eye=[0.3, 0, 0.6], target=[-0.1, 0, 0.1])
		pose = sapien_utils.look_at(eye=[1, 0.0, 0.6], target=[-0.2, 0.0, 0.2])
		return [CameraConfig("base_camera", pose, 300, 300, np.pi / 2, 0.01, 100)]

		@property
		def _default_human_render_camera_configs(self):
		pose = sapien_utils.look_at([0.6, 0.7, 0.6], [0.0, 0.0, 0.35])
		pose = sapien_utils.look_at([1, 0.0, 0.6], [-0.2, 0.0, 0.2])
		return CameraConfig("render_camera", pose, 300,300, 1, 0.01, 100)

		def _load_scene(self, options: dict):
		@@ -62,11 +64,11 @@ class PlaceTwoCubeEnv(BaseEnv):
		self.table_scene.build()

		# Create two cubes
		self.cubeA = actors.build_cube(
		self.scene, half_size=0.02, color=[1, 0, 0, 1], name="cubeA"
		self.red_cube = actors.build_cube(
		self.scene, half_size=0.02, color=[1, 0, 0, 1], name="red_cube"
		)
		self.cubeB = actors.build_cube(
		self.scene, half_size=0.02, color=[0, 1, 0, 1], name="cubeB"
		self.green_cube = actors.build_cube(
		self.scene, half_size=0.02, color=[0, 1, 0, 1], name="green_cube"
		)

		# Create two target areas
		@@ -96,48 +98,48 @@ class PlaceTwoCubeEnv(BaseEnv):
		# Place cubes at random positions
		xyz = torch.zeros((b, 3))
		xyz[:, 2] = 0.02
		region = [[self.workspace_x[0], self.workspace_y[0]], [self.workspace_x[1], -0.02]]
		region =self.region
		sampler = randomization.UniformPlacementSampler(bounds=region, batch_size=b)
		radius = torch.linalg.norm(torch.tensor([0.02, 0.02])) + 0.02

		cubeA_xy = sampler.sample(radius, 100)
		cubeB_xy = sampler.sample(radius, 100, verbose=False)
		red_cube_xy = sampler.sample(radius, 100)
		green_cube_xy = sampler.sample(radius, 100, verbose=False)

		# Set initial positions for cubes
		xyz[:, :2] = cubeA_xy
		xyz[:, :2] = red_cube_xy
		qs = randomization.random_quaternions(b, lock_x=True, lock_y=True, lock_z=False)
		self.cubeA.set_pose(Pose.create_from_pq(p=xyz.clone(), q=qs))
		self.red_cube.set_pose(Pose.create_from_pq(p=xyz.clone(), q=qs))

		xyz[:, :2] = cubeB_xy
		xyz[:, :2] = green_cube_xy
		qs = randomization.random_quaternions(b, lock_x=True, lock_y=True, lock_z=False)
		self.cubeB.set_pose(Pose.create_from_pq(p=xyz.clone(), q=qs))
		self.green_cube.set_pose(Pose.create_from_pq(p=xyz.clone(), q=qs))

		# Set fixed positions for target areas
		goalA_pos = torch.tensor([0.1, 0.1, 0.0])
		goalB_pos = torch.tensor([-0.1, 0.1, 0.0])
		left_target_position = torch.tensor([0.08, -0.1, 0.0])
		right_target_position = torch.tensor([0.08, 0.1, 0.0])

		self.goal_region_A.set_pose(Pose.create_from_pq(
		p=goalA_pos,
		p=left_target_position,
		q=euler2quat(0, np.pi / 2, 0),
		))
		self.goal_region_B.set_pose(Pose.create_from_pq(
		p=goalB_pos,
		p=right_target_position,
		q=euler2quat(0, np.pi / 2, 0),
		))

		def is_cube_in_goal(self, cube_pos, goal_pos):
		distance = torch.norm(cube_pos[..., :2] - goal_pos[...,:2], dim=-1)
		def is_cube_in_goal(self, cube_position, goal_position):
		distance = torch.norm(cube_position[..., :2] - goal_position[...,:2], dim=-1)
		return distance <= self.goal_radius

		def _get_obs_extra(self, info: Dict):
		assert "state" in self.obs_mode
		obs = dict(
		cubeA_pos=info["cubeA_pos"],
		cubeB_pos=info["cubeB_pos"],
		is_cubeA_grasped=info["is_cubeA_grasped"],
		is_cubeB_grasped=info["is_cubeB_grasped"],
		goalA_pos=self.goal_region_A.pose.p,
		goalB_pos=self.goal_region_B.pose.p,
		red_cube_position=info["red_cube_position"],
		green_cube_position=info["green_cube_position"],
		is_red_cube_grasped=info["is_red_cube_grasped"],
		is_green_cube_grasped=info["is_green_cube_grasped"],
		left_target_position=self.goal_region_A.pose.p,
		right_target_position=self.goal_region_B.pose.p,
		)
		return obs

		@@ -153,29 +155,29 @@ class PlaceTwoCubeEnv(BaseEnv):

		def task_fail(self, info: Dict):
		# if cube position is out of workspace return true
		for cube in ["cubeA", "cubeB"]:
		if info[f"{cube}_pos"][0] < self.workspace_x[0] or info[f"{cube}_pos"][0] > self.workspace_x[1]:
		for cube in ["red_cube", "green_cube"]:
		if info[f"{cube}_position"][0] < self.workspace_x[0] or info[f"{cube}_position"][0] > self.workspace_x[1]:
		return True
		if info[f"{cube}_pos"][1] < self.workspace_y[0] or info[f"{cube}_pos"][1] > self.workspace_y[1]:
		if info[f"{cube}_position"][1] < self.workspace_y[0] or info[f"{cube}_position"][1] > self.workspace_y[1]:
		return True
		if info[f"{cube}_pos"][2] < 0:
		if info[f"{cube}_position"][2] < 0:
		return True
		return False

		def evaluate(self):
		pos_A = self.cubeA.pose.p
		pos_B = self.cubeB.pose.p
		pos_A = self.red_cube.pose.p
		pos_B = self.green_cube.pose.p

		is_cubeA_grasped = self.agent.is_grasping(self.cubeA)
		is_cubeB_grasped = self.agent.is_grasping(self.cubeB)
		is_red_cube_grasped = self.agent.is_grasping(self.red_cube)
		is_green_cube_grasped = self.agent.is_grasping(self.green_cube)

		is_cubeA_in_goal = self.is_cube_in_goal(pos_A, self.goal_region_A.pose.p)
		is_cubeB_in_goal = self.is_cube_in_goal(pos_B, self.goal_region_B.pose.p)
		is_red_cube_in_goal = self.is_cube_in_goal(pos_A, self.goal_region_A.pose.p)
		is_green_cube_in_goal = self.is_cube_in_goal(pos_B, self.goal_region_B.pose.p)

		stage0_success = is_cubeA_grasped
		stage1_success = is_cubeA_in_goal & (~is_cubeA_grasped)
		stage2_success = is_cubeA_in_goal & is_cubeB_grasped
		stage3_success = is_cubeA_in_goal & is_cubeB_in_goal & (~is_cubeB_grasped) & (~is_cubeA_grasped)
		stage0_success = is_red_cube_grasped
		stage1_success = is_red_cube_in_goal & (~is_red_cube_grasped)
		stage2_success = is_red_cube_in_goal & is_green_cube_grasped
		stage3_success = is_red_cube_in_goal & is_green_cube_in_goal & (~is_green_cube_grasped) & (~is_red_cube_grasped)



		@@ -184,14 +186,14 @@ class PlaceTwoCubeEnv(BaseEnv):


		info = {
		"goalA_pos":self.goal_region_A.pose.p,
		"goalB_pos":self.goal_region_B.pose.p,
		"is_cubeA_grasped": is_cubeA_grasped,
		"is_cubeB_grasped": is_cubeB_grasped,
		"cubeA_pos": pos_A,
		"cubeB_pos": pos_B,
		"is_cubeA_in_goal": is_cubeA_in_goal,
		"is_cubeB_in_goal": is_cubeB_in_goal,
		"left_target_position":self.goal_region_A.pose.p,
		"right_target_position":self.goal_region_B.pose.p,
		"is_red_cube_grasped": is_red_cube_grasped,
		"is_green_cube_grasped": is_green_cube_grasped,
		"red_cube_position": pos_A,
		"green_cube_position": pos_B,
		"is_red_cube_in_goal": is_red_cube_in_goal,
		"is_green_cube_in_goal": is_green_cube_in_goal,
		"stage0_success": stage0_success.bool(),
		"stage1_success": stage1_success.bool(),
		"stage2_success": stage2_success.bool(),

vagen/env/primitive_skill/maniskill/env/put_apple_in_drawer.py

+12 −11

Original line number	Diff line number	Diff line
		@@ -22,6 +22,7 @@ class PutAppleInDrawerEnv(BaseEnv):
		SUPPORTED_ROBOTS = ["panda", "xmate3_robotiq", "fetch"]
		agent: Union[Panda, Xmate3Robotiq, Fetch]
		vlm_info_keys=["drawer_open_value"]
		state_keys=["apple_position", "drawer_position"]
		# Asset configuration and constants
		DRAWER_ASSET_ID = "partnet_mobility_cabinet"
		handle_types = ["prismatic"] # We are interested in prismatic joints (drawers)
		@@ -51,12 +52,12 @@ class PutAppleInDrawerEnv(BaseEnv):

		@property
		def _default_sensor_configs(self):
		pose = sapien_utils.look_at(eye=[-0.5, 0, 0.6], target=[-0.5, 0, 0.1])
		return [CameraConfig("base_camera", pose, 128, 128, np.pi / 2, 0.01, 100)]
		pose = sapien_utils.look_at(eye=[1, 0.0, 0.6], target=[-0.2, 0.0, 0.2])
		return [CameraConfig("base_camera", pose, 300, 300, np.pi / 2, 0.01, 100)]

		@property
		def _default_human_render_camera_configs(self):
		pose = sapien_utils.look_at([1.0, 0.3, 1.0], [-0.5, -0.2, 0.35])
		pose = sapien_utils.look_at([1, 0.0, 0.6], [-0.2, 0.0, 0.2])
		return CameraConfig("render_camera", pose, 300,300, 1, 0.01, 100)

		def instruction(self):
		@@ -141,16 +142,16 @@ class PutAppleInDrawerEnv(BaseEnv):
		info= {}
		for name in self.object_list:
		info[f"is_{name}_grasped"] = self.agent.is_grasping(self.object_list[name])[0]
		info[f"{name}_pos"] = self.object_list[name].pose.p[0]
		info[f"{name}_position"] = self.object_list[name].pose.p[0]

		info["stage"] = self.cur_stage
		info["gripper_pos"] = self.agent.tcp.pose.p[0]
		info["gripper_position"] = self.agent.tcp.pose.p[0]

		drawer_link = self.drawer_joint.get_child_link()
		info["drawer_handle_pos"] = drawer_link.pose.p.to(self.device)[0] + np.array([0, 0.37, -0.3])
		info["drawer_pos"] = drawer_link.pose.p.to(self.device)[0] + np.array([0, 0.2, -0.3])
		info["drawer_handle_position"] = drawer_link.pose.p.to(self.device)[0] + np.array([0, 0.37, -0.3])
		info["drawer_position"] = drawer_link.pose.p.to(self.device)[0] + np.array([0, 0.2, -0.3])

		info["drawer_open_value"] = (info["drawer_handle_pos"][1] - np.array([0.5, -0.63, 0.2]))[1]
		info["drawer_open_value"] = (info["drawer_handle_position"][1] - np.array([0.5, -0.63, 0.2]))[1]


		return info
		@@ -163,7 +164,7 @@ class PutAppleInDrawerEnv(BaseEnv):

		def stage1_success(info):
		# Check if apple is in drawer
		abs_diff_xy = torch.abs(info["apple_pos"][:2] - info["drawer_pos"][:2])
		abs_diff_xy = torch.abs(info["apple_pos"][:2] - info["drawer_position"][:2])

		# Check if x and y differences are within tolerance
		within_xy = (abs_diff_xy <= 0.1).all(dim=-1)
		@@ -177,7 +178,7 @@ class PutAppleInDrawerEnv(BaseEnv):

		def stage2_success(info):
		# Check if apple is in drawer
		abs_diff_xy = torch.abs(info["apple_pos"][:2] - info["drawer_pos"][:2])
		abs_diff_xy = torch.abs(info["apple_pos"][:2] - info["drawer_position"][:2])

		# Check if x and y differences are within tolerance
		within_xy = (abs_diff_xy <= 0.1).all(dim=-1)
		@@ -203,7 +204,7 @@ class PutAppleInDrawerEnv(BaseEnv):
		info = self.get_info()
		obs = []
		for name in self.object_list:
		obs += info[f"{name}_pos"].flatten().tolist()
		obs += info[f"{name}_position"].flatten().tolist()

		for name in self.object_list:
		obs += info[f"is_{name}_grasped"].flatten().tolist()

vagen/env/primitive_skill/maniskill/env/stack_three_cube.py

+13 −13

Original line number	Diff line number	Diff line
		@@ -23,6 +23,7 @@ class StackThreeCubeEnv(BaseEnv):
		agent: Union[Panda, Xmate3Robotiq, Fetch]
		skill_config=None
		vlm_info_keys=['cube_size']
		state_keys=["red_cube_position", "green_cube_position", "purple_cube_position"]

		def __init__(self, stage=0,args, robot_uids="panda", robot_init_qpos_noise=0.02, *kwargs):
		self.stage=stage
		@@ -30,8 +31,9 @@ class StackThreeCubeEnv(BaseEnv):
		self.cube_size = 0.04

		self.workspace_x=[-0.10, 0.15]
		self.workspace_y=[-0.15, 0.15]
		self.workspace_z=[0.01, 0.15]
		self.workspace_y=[-0.2, 0.2]
		self.workspace_z=[0.01, 0.2]

		self.robot_init_qpos_noise = robot_init_qpos_noise
		self.reward_components = ["success","afford"]
		super().__init__(args, robot_uids=robot_uids, *kwargs)
		@@ -46,15 +48,13 @@ class StackThreeCubeEnv(BaseEnv):

		@property
		def _default_sensor_configs(self):
		pose = sapien_utils.look_at(eye=[0.3, 0, 0.6], target=[-0.1, 0, 0.1])
		pose = sapien_utils.look_at(eye=[1, 0.0, 0.6], target=[-0.2, 0.0, 0.2])
		return [CameraConfig("base_camera", pose, 300, 300, np.pi / 2, 0.01, 100)]

		@property
		def _default_human_render_camera_configs(self):
		pose = sapien_utils.look_at([0.8, 0.8, 0.6], [0.14, +0.08, 0.12])
		return CameraConfig(
		"render_camera", pose=pose, width=300, height=300, fov=0.63, near=0.01, far=100
		)
		pose = sapien_utils.look_at([1, 0.0, 0.6], [-0.2, 0.0, 0.2])
		return CameraConfig("render_camera", pose, 300,300, 1, 0.01, 100)

		def instruction(self):
		return "Please stack the red cube on top of the green cube, and then stack purple cube on top of the red cube."
		@@ -124,9 +124,9 @@ class StackThreeCubeEnv(BaseEnv):
		info = {}
		for name in self.object_list:
		info[f"is_{name}_grasped"] = self.agent.is_grasping(self.object_list[name])[0]
		info[f"{name}_pos"] = self.object_list[name].pose.p[0]
		info["cube_size"]=torch.ones_like(info["red_cube_pos"])*0.04
		info["gripper_pos"] = self.agent.tcp.pose.p[0]
		info[f"{name}_position"] = self.object_list[name].pose.p[0]
		info["cube_size"]=torch.ones_like(info["red_cube_position"])*0.04
		info["gripper_position"] = self.agent.tcp.pose.p[0]
		return info

		def evaluate(self):
		@@ -137,7 +137,7 @@ class StackThreeCubeEnv(BaseEnv):

		def stage1_success(info):
		red_not_grasped = ~info["is_red_cube_grasped"]
		red_on_green = (torch.linalg.norm(info["red_cube_pos"][:2] - info["green_cube_pos"][:2]) < self.cube_size/2) and (info["red_cube_pos"][2] > (info["green_cube_pos"][2] + self.cube_size/2))
		red_on_green = (torch.linalg.norm(info["red_cube_position"][:2] - info["green_cube_position"][:2]) < self.cube_size/2) and (info["red_cube_position"][2] > (info["green_cube_position"][2] + self.cube_size/2))
		return (red_on_green and red_not_grasped)

		def stage2_success(info):
		@@ -145,7 +145,7 @@ class StackThreeCubeEnv(BaseEnv):

		def stage3_success(info):
		purple_not_grasped = ~info["is_purple_cube_grasped"]
		purple_on_red = (torch.linalg.norm(info["purple_cube_pos"][:2] - info["red_cube_pos"][:2]) < self.cube_size/2) and (info["purple_cube_pos"][2] > (info["red_cube_pos"][2] + self.cube_size/2))
		purple_on_red = (torch.linalg.norm(info["purple_cube_position"][:2] - info["red_cube_position"][:2]) < self.cube_size/2) and (info["purple_cube_position"][2] > (info["red_cube_position"][2] + self.cube_size/2))

		return purple_on_red and purple_not_grasped and stage1_success(info)

		@@ -161,7 +161,7 @@ class StackThreeCubeEnv(BaseEnv):
		info = self.get_info()
		obs = []
		for name in self.object_list:
		obs += info[f"{name}_pos"].flatten().tolist()
		obs += info[f"{name}_position"].flatten().tolist()

		for name in self.object_list:
		obs += info[f"is_{name}_grasped"].flatten().tolist()

vagen/env/primitive_skill/maniskill/utils.py

+2 −2

Original line number	Diff line number	Diff line
		@@ -38,7 +38,7 @@ def handle_info(info,mask_success=False,env=None):
		info.pop(k)

		for k, v in info.items():
		if k.endswith('_pos'):
		if k.endswith('_position'):
		# Convert position arrays to integer tuples in cm
		obj_positions[k] = tuple(np.round(v * 1000, 0).astype(int).tolist())
		elif k.endswith('_value'):
		@@ -70,7 +70,7 @@ def handle_info(info,mask_success=False,env=None):
		final_info.update(other_info)

		return {
		'obj_positions': obj_positions,
		'obj_positions': list(obj_positions.values()),
		'other_info': str(final_info)
		}

Admin message