Robocasa dev

robosuite/controllers/parts/mobile_base/joint_vel.py

@@ -206,14 +206,33 @@ def set_goal(self, action, set_qpos=None):
         curr_theta = T.mat2euler(curr_ori)[2]  # np.arctan2(curr_pos[1], curr_pos[0])
         theta = curr_theta - init_theta
 
-        # input raw base action is delta relative to current pose of base
-        # controller expects deltas relative to initial pose of base at start of episode
-        # transform deltas from current base pose coordinates to initial base pose coordinates
-        x, y = action[0:2]
-
-        # do the reverse of theta rotation
-        action[0] = x * np.cos(theta) + y * np.sin(theta)
-        action[1] = -x * np.sin(theta) + y * np.cos(theta)
+        # Detect the axis for the forward joint and dynamically reorder action accordingly.
+        # This is needed because previous versions of the mobile base xml had different forward
+        # axis definitions. In order to maintain backwards compatibility with previous datasets
+        # we dynamically detect the forward joint axis and reorder the action accordingly.
+        forward_jnt = None
+        forward_jnt_axis = None
+        for jnt in self.joint_names:
+            if "joint_mobile_forward" in jnt:
+                forward_jnt = jnt
+                forward_jnt_axis = self.sim.model.jnt_axis[self.sim.model.joint_name2id(jnt)]
+                break
+
+        # reorder action if forward axis is y axis
+        if forward_jnt and (forward_jnt_axis == np.array([0, 1, 0])).all():
+            action = np.copy([action[i] for i in [1, 0, 2]])
+
+            x, y = action[0:2]
+            # do the reverse of theta rotation
+            action[0] = x * np.cos(theta) + y * np.sin(theta)
+            action[1] = -x * np.sin(theta) + y * np.cos(theta)
+        else:
+            # input raw base action is delta relative to current pose of base
+            # controller expects deltas relative to initial pose of base at start of episode
+            # transform deltas from current base pose coordinates to initial base pose coordinates
+            x, y = action[0:2]
+            action[0] = x * np.cos(theta) - y * np.sin(theta)
+            action[1] = x * np.sin(theta) + y * np.cos(theta)
 
         self.goal_qvel = action
         if self.interpolator is not None:

robosuite/devices/device.py

@@ -21,8 +21,7 @@ def __init__(self, env):
                             using this device.
         """
         self.env = env
-        self.all_robot_arms = [robot.arms for robot in self.env.robots]
-        self.num_robots = len(self.all_robot_arms)
+        self._all_robot_arms = None
 
     def _reset_internal_state(self):
         """
@@ -36,6 +35,16 @@ def _reset_internal_state(self):
         self._prev_target = {arm: None for arm in self.all_robot_arms[self.active_robot]}
         self._prev_torso_target = None
 
+    @property
+    def all_robot_arms(self):
+        if self._all_robot_arms is None:
+            self._all_robot_arms = [robot.arms for robot in self.env.robots]
+        return self._all_robot_arms
+
+    @property
+    def num_robots(self):
+        return len(self.all_robot_arms)
+
     @property
     def active_arm(self):
         return self.all_robot_arms[self.active_robot][self.active_arm_index]

robosuite/devices/keyboard.py

@@ -23,7 +23,6 @@ def __init__(self, env, pos_sensitivity=1.0, rot_sensitivity=1.0):
         super().__init__(env)
 
         self._display_controls()
-        self._reset_internal_state()
 
         self._reset_state = 0
         self._enabled = False

robosuite/environments/base.py

@@ -82,6 +82,8 @@ class MujocoEnv(metaclass=EnvMeta):
         ignore_done (bool): True if never terminating the environment (ignore @horizon).
         hard_reset (bool): If True, re-loads model, sim, and render object upon a reset call, else,
             only calls sim.reset and resets all robosuite-internal variables
+        load_model_on_init (bool): If True, load and initialize the model and renderer in __init__ constructor,
+            else, initialize these components in the first call to reset()
         renderer (str): string for the renderer to use
         renderer_config (dict): dictionary for the renderer configurations
         seed (int): environment seed. Default is None, where environment is unseeded, ie. random
@@ -102,6 +104,7 @@ def __init__(
         horizon=1000,
         ignore_done=False,
         hard_reset=True,
+        load_model_on_init=True,
         renderer="mjviewer",
         renderer_config=None,
         seed=None,
@@ -143,31 +146,44 @@ def __init__(
 
         self._ep_meta = {}
 
-        # Load the model
-        self._load_model()
+        self.load_model_on_init = load_model_on_init
 
-        # Initialize the simulation
-        self._initialize_sim()
+        # variable to keep track of whether the env has been fully initialized
+        self._env_is_initialized = False
 
-        # initializes the rendering
-        self.initialize_renderer()
+        if self.load_model_on_init:
+            # Load the model
+            self._load_model()
 
-        # the variables will be set later.
-        # need to set to None, in case these variables are referenced before being set
-        self.viewer = None
-        self.viewer_get_obs = None
+            # Initialize the simulation
+            self._initialize_sim()
 
-        # Run all further internal (re-)initialization required
-        self._reset_internal()
+            # initializes the rendering
+            self.initialize_renderer()
 
-        # Load observables
-        if hasattr(self.viewer, "_setup_observables"):
-            self._observables = self.viewer._setup_observables()
-        else:
-            self._observables = self._setup_observables()
+            # the variables will be set later.
+            # need to set to None, in case these variables are referenced before being set
+            self.viewer = None
+            self.viewer_get_obs = None
 
-        # check if viewer has get observations method and set a flag for future use.
-        self.viewer_get_obs = hasattr(self.viewer, "_get_observations")
+            # Run all further internal (re-)initialization required
+            self._reset_internal()
+
+            # Load observables
+            if hasattr(self.viewer, "_setup_observables"):
+                self._observables = self.viewer._setup_observables()
+            else:
+                self._observables = self._setup_observables()
+
+            # check if viewer has get observations method and set a flag for future use.
+            self.viewer_get_obs = hasattr(self.viewer, "_get_observations")
+            self._env_is_initialized = True
+        else:
+            # the variables will be set later.
+            # need to set to None, in case these variables are referenced before being set
+            self.sim = None
+            self.viewer = None
+            self.viewer_get_obs = None
 
     def initialize_renderer(self):
         self.renderer = self.renderer.lower()
@@ -271,7 +287,7 @@ def reset(self):
         if self.renderer == "mjviewer":
             self._destroy_viewer()
 
-        if self.hard_reset and not self.deterministic_reset:
+        if (self.sim is None) or (self.hard_reset and not self.deterministic_reset):
             if self.renderer == "mujoco":
                 self._destroy_viewer()
                 self._destroy_sim()
@@ -281,9 +297,33 @@ def reset(self):
         else:
             self.sim.reset()
 
-        # Reset necessary robosuite-centric variables
-        self._reset_internal()
-        self.sim.forward()
+        if self._env_is_initialized is True:
+            # Reset necessary robosuite-centric variables
+            self._reset_internal()
+            self.sim.forward()
+        else:
+            # initializes the rendering
+            self.initialize_renderer()
+
+            # the variables will be set later.
+            # need to set to None, in case these variables are referenced before being set
+            self.viewer = None
+            self.viewer_get_obs = None
+
+            # Run all further internal (re-)initialization required
+            self._reset_internal()
+            self.sim.forward()
+
+            # Load observables
+            if hasattr(self.viewer, "_setup_observables"):
+                self._observables = self.viewer._setup_observables()
+            else:
+                self._observables = self._setup_observables()
+
+            # check if viewer has get observations method and set a flag for future use.
+            self.viewer_get_obs = hasattr(self.viewer, "_get_observations")
+            self._env_is_initialized = True
+
         # Setup observables, reloading if
         self._obs_cache = {}
         self._reset_observables()

robosuite/environments/manipulation/manipulation_env.py

@@ -84,6 +84,9 @@ class ManipulationEnv(RobotEnv):
         hard_reset (bool): If True, re-loads model, sim, and render object upon a reset call, else,
             only calls sim.reset and resets all robosuite-internal variables
 
+        load_model_on_init (bool): If True, load and initialize the model and renderer in __init__ constructor,
+            else, initialize these components in the first call to reset()
+
         camera_names (str or list of str): name of camera to be rendered. Should either be single str if
             same name is to be used for all cameras' rendering or else it should be a list of cameras to render.
 
@@ -144,6 +147,7 @@ def __init__(
         horizon=1000,
         ignore_done=False,
         hard_reset=True,
+        load_model_on_init=True,
         camera_names="agentview",
         camera_heights=256,
         camera_widths=256,
@@ -187,6 +191,7 @@ def __init__(
             horizon=horizon,
             ignore_done=ignore_done,
             hard_reset=hard_reset,
+            load_model_on_init=load_model_on_init,
             camera_names=camera_names,
             camera_heights=camera_heights,
             camera_widths=camera_widths,

robosuite/environments/robot_env.py

@@ -78,6 +78,9 @@ class RobotEnv(MujocoEnv):
         hard_reset (bool): If True, re-loads model, sim, and render object upon a reset call, else,
             only calls sim.reset and resets all robosuite-internal variables
 
+        load_model_on_init (bool): If True, load and initialize the model and renderer in __init__ constructor,
+            else, initialize these components in the first call to reset()
+
         camera_names (str or list of str): name of camera to be rendered. Should either be single str if
             same name is to be used for all cameras' rendering or else it should be a list of cameras to render.
 
@@ -139,6 +142,7 @@ def __init__(
         horizon=1000,
         ignore_done=False,
         hard_reset=True,
+        load_model_on_init=True,
         camera_names="agentview",
         camera_heights=256,
         camera_widths=256,
@@ -230,6 +234,7 @@ def __init__(
             horizon=horizon,
             ignore_done=ignore_done,
             hard_reset=hard_reset,
+            load_model_on_init=load_model_on_init,
             renderer=renderer,
             renderer_config=renderer_config,
             seed=seed,

robosuite/models/tasks/task.py

@@ -51,6 +51,8 @@ class Task(MujocoWorldBase):
 
         mujoco_objects (None or MujocoObject or list of MujocoObject): a list of MJCF models of physical objects
 
+        enable_multiccd (bool) whether to set the multiccd flag in MuJoCo. False by default
+
     Raises:
         AssertionError: [Invalid input object type]
     """
@@ -60,8 +62,10 @@ def __init__(
         mujoco_arena,
         mujoco_robots,
         mujoco_objects=None,
+        enable_multiccd=False,
+        enable_sleeping_islands=False,
     ):
-        super().__init__()
+        super().__init__(enable_multiccd=enable_multiccd, enable_sleeping_islands=enable_sleeping_islands)
 
         # Store references to all models
         self.mujoco_arena = mujoco_arena

robosuite/models/world.py

@@ -1,3 +1,5 @@
+import xml.etree.ElementTree as ET
+
 import robosuite.macros as macros
 from robosuite.models.base import MujocoXML
 from robosuite.utils.mjcf_utils import convert_to_string, find_elements, xml_path_completion
@@ -6,8 +8,18 @@
 class MujocoWorldBase(MujocoXML):
     """Base class to inherit all mujoco worlds from."""
 
-    def __init__(self):
+    def __init__(self, enable_multiccd=False, enable_sleeping_islands=False):
         super().__init__(xml_path_completion("base.xml"))
         # Modify the simulation timestep to be the requested value
         options = find_elements(root=self.root, tags="option", attribs=None, return_first=True)
         options.set("timestep", convert_to_string(macros.SIMULATION_TIMESTEP))
+        self.enable_multiccd = enable_multiccd
+        self.enable_sleeping_islands = enable_sleeping_islands
+        if self.enable_multiccd:
+            multiccd_elem = ET.fromstring("""<option> <flag multiccd="enable"/> </option>""")
+            mujoco_elem = find_elements(self.root, "mujoco")
+            mujoco_elem.insert(0, multiccd_elem)
+        if self.enable_sleeping_islands:
+            sleeping_elem = ET.fromstring("""<option> <flag sleep="enable"/> </option>""")
+            mujoco_elem = find_elements(self.root, "mujoco")
+            mujoco_elem.insert(0, sleeping_elem)

robosuite/wrappers/data_collection_wrapper.py

@@ -14,7 +14,7 @@
 
 
 class DataCollectionWrapper(Wrapper):
-    def __init__(self, env, directory, collect_freq=1, flush_freq=100):
+    def __init__(self, env, directory, collect_freq=1, flush_freq=100, use_env_xml_for_reset=False):
         """
         Initializes the data collection wrapper.
 
@@ -23,11 +23,14 @@ def __init__(self, env, directory, collect_freq=1, flush_freq=100):
             directory (str): Where to store collected data.
             collect_freq (int): How often to save simulation state, in terms of environment steps.
             flush_freq (int): How frequently to dump data to disk, in terms of environment steps.
+            use_env_xml_for_reset (bool): Whether to use the robosuite env XML string or the xml
+                                          string from self.sim for resetting the environment.
         """
         super().__init__(env)
 
         # the base directory for all logging
         self.directory = directory
+        self.use_env_xml_for_reset = use_env_xml_for_reset
 
         # in-memory cache for simulation states and action info
         self.states = []
@@ -69,10 +72,16 @@ def _start_new_episode(self):
 
         # save the task instance (will be saved on the first env interaction)
 
-        # NOTE: was previously self.env.model.get_xml(). Was causing the following issue in rare cases:
+        # NOTE: was originally set to self.env.model.get_xml()
+        # That was causing the following issue in rare cases:
         # ValueError: Error: eigenvalues of mesh inertia violate A + B >= C
-        # switching to self.env.sim.model.get_xml() does not create this issue
-        self._current_task_instance_xml = self.env.sim.model.get_xml()
+        # then, switched to self.env.sim.model.get_xml() which does not create this issue
+        # however, that leads to subtle changes in the physics, such as fixture doors being harder to close
+        # so, in order to address both issues, added an flag to choose between the two methods
+        if self.use_env_xml_for_reset:
+            self._current_task_instance_xml = self.env.model.get_xml()
+        else:
+            self._current_task_instance_xml = self.env.sim.model.get_xml()
         self._current_task_instance_state = np.array(self.env.sim.get_state().flatten())
 
         # trick for ensuring that we can play MuJoCo demonstrations back

robosuite/wrappers/visualization_wrapper.py

@@ -75,9 +75,10 @@ def __init__(self, env, indicator_configs=None):
 
         # Conduct a (hard) reset to make sure visualization changes propagate
         reset_mode = self.env.hard_reset
-        self.env.hard_reset = True
-        self.reset()
-        self.env.hard_reset = reset_mode
+        if self.env.hard_reset is not True:
+            self.env.hard_reset = True
+            self.reset()
+            self.env.hard_reset = reset_mode
 
     def get_indicator_names(self):
         """