From 602b7a4e680fbb9e7e4fca398b6421e675660cb6 Mon Sep 17 00:00:00 2001 From: arth-shukla <44527268+arth-shukla@users.noreply.github.com> Date: Thu, 29 Feb 2024 23:59:07 -0800 Subject: [PATCH] Fix: merge actors by _objs --- mani_skill2/envs/scenes/base_env.py | 2 + .../envs/scenes/tasks/sequential_task.py | 78 +++++++++---------- mani_skill2/utils/building/actors.py | 3 + 3 files changed, 44 insertions(+), 39 deletions(-) diff --git a/mani_skill2/envs/scenes/base_env.py b/mani_skill2/envs/scenes/base_env.py index d69c32417..95091b49c 100644 --- a/mani_skill2/envs/scenes/base_env.py +++ b/mani_skill2/envs/scenes/base_env.py @@ -63,6 +63,8 @@ def reset(self, seed=None, options=None): self._set_episode_rng(seed) if options is None: options = dict(reconfigure=False) + if "reconfigure" not in options: + options["reconfigure"] = False if not self.fixed_scene: options["reconfigure"] = True if options["reconfigure"]: diff --git a/mani_skill2/envs/scenes/tasks/sequential_task.py b/mani_skill2/envs/scenes/tasks/sequential_task.py index feffc68a6..1ec28b04f 100644 --- a/mani_skill2/envs/scenes/tasks/sequential_task.py +++ b/mani_skill2/envs/scenes/tasks/sequential_task.py @@ -1,5 +1,5 @@ from collections import OrderedDict -from typing import Any, Dict, List, Tuple +from typing import Any, Dict, List, Tuple, Union import numpy as np import sapien @@ -89,14 +89,11 @@ def __init__( ): if "num_envs" in kwargs: - assert ( - len(task_plans) == kwargs["num_envs"], - f"GPU sim requires equal number of task_plans ({len(task_plans)}) and parallel_envs ({kwargs['num_envs']})", - ) + assert len(task_plans) == kwargs["num_envs"] or len(task_plans) == 1, \ + f"GPU sim requires equal number of task_plans ({len(task_plans)}) and parallel_envs ({kwargs['num_envs']})" else: - assert ( - len(task_plans) == 1 - ), f"CPU sim only supports one task plan, not {(len(task_plans))}" + assert len(task_plans) == 1, \ + f"CPU sim only supports one task plan, not {(len(task_plans))}" self.base_task_plans = task_plans self.robot_init_qpos_noise = robot_init_qpos_noise @@ -137,16 +134,9 @@ def process_task_plan(self): if isinstance(subtask0, PickSubtask): parallel_subtasks: List[PickSubtask] merged_obj_name = f"obj_{subtask_num}" - - self.subtask_objs.append( - Actor.merge( - [ - self._get_actor(subtask.obj_id) - for subtask in parallel_subtasks - ], - name=merged_obj_name, - ) - ) + self.subtask_objs.append(self._create_merged_actor_from_subtasks( + parallel_subtasks, name=merged_obj_name + )) self.subtask_goals.append(None) self.task_plan.append(PickSubtask(obj_id=merged_obj_name)) @@ -155,26 +145,13 @@ def process_task_plan(self): parallel_subtasks: List[PlaceSubtask] merged_obj_name = f"obj_{subtask_num}" merged_goal_name = f"goal_{subtask_num}" - - self.subtask_objs.append( - Actor.merge( - [ - self._get_actor(subtask.obj_id) - for subtask in parallel_subtasks - ], - name=merged_obj_name, - ) - ) + self.subtask_objs.append(self._create_merged_actor_from_subtasks( + parallel_subtasks, name=merged_obj_name + )) self.subtask_goals.append( - Actor.merge( - [ - self._make_goal( - pos=subtask.goal_pos, - radius=self.place_cfg.obj_goal_thresh, - name=f"goal_{subtask.uid}", - ) - for subtask in parallel_subtasks - ], + self._make_goal( + pos=torch.tensor(subtask.goal_pos for subtask in parallel_subtasks), + radius=self.place_cfg.obj_goal_thresh, name=merged_goal_name, ) ) @@ -207,12 +184,31 @@ def process_task_plan(self): def _get_actor(self, actor_id: str): return self.scene_builder.movable_objects_by_id[actor_id] + + def _create_merged_actor_from_subtasks( + self, + parallel_subtasks: Union[List[PickSubtask], List[PlaceSubtask]], + name: str = None + ): + merged_obj = Actor._create_from_entities( + [ + self._get_actor(subtask.obj_id)._objs[i] + for i, subtask in enumerate(parallel_subtasks) + ], + scene=self._scene, + scene_mask=np.ones(self.num_envs, dtype=bool) + ) + if name is not None: + merged_obj.name = name, + return merged_obj + def _make_goal( self, - pos: Tuple[float, float, float] = None, + pos: Union[Tuple[float, float, float], List[Tuple[float, float, float]]] = None, radius=0.15, name="goal_site", + scene_mask=None, ): goal = build_sphere( self._scene, @@ -221,9 +217,13 @@ def _make_goal( name=name, body_type="kinematic", add_collision=False, + scene_mask=scene_mask, ) if pos is not None: - goal.set_pose(sapien.Pose(p=pos)) + if len(pos) == self.num_envs: + goal.set_pose(Pose.create_from_pq(p=pos)) + else: + goal.set_pose(sapien.Pose(p=pos)) self._hidden_objects.append(goal) return goal diff --git a/mani_skill2/utils/building/actors.py b/mani_skill2/utils/building/actors.py index 04bb8337d..f2876ad90 100644 --- a/mani_skill2/utils/building/actors.py +++ b/mani_skill2/utils/building/actors.py @@ -84,8 +84,11 @@ def build_sphere( name: str, body_type: str = "dynamic", add_collision: bool = True, + scene_mask=None, ): builder = scene.create_actor_builder() + if scene_mask is not None: + builder.set_scene_mask(scene_mask) if add_collision: builder.add_sphere_collision( radius=radius,