v0.6.0 Alpha Release: Sapien 3 integration (haosulab#175)

* version bump * copy code over * refactor xmate3 * Update xmate3.py * Changed Replay Trajectory to add options to record rewards as well as specify reward mode (haosulab#177) * Added reward mode argument and record rewards argument to replay trajectory * Fixed an error --------- Co-authored-by: Arnav G <[email protected]> Co-authored-by: Stone Tao <[email protected]> * merge in refactor and updates for utils * merge in changes for trajectory * merge new examples * merge new env code, refactors + new scene builder api * merge new agents code * merge code for new assets, floors etc. * Update __init__.py * Delete README.md --------- Co-authored-by: Arnav G <[email protected]> Co-authored-by: Arnav G <[email protected]>
itsygao · Jan 30, 2024 · 4d59009 · 4d59009
1 parent 8e1e482
commit 4d59009
Show file tree

Hide file tree

Showing 154 changed files with 16,430 additions and 1,615 deletions.
diff --git a/.gitignore b/.gitignore
@@ -38,4 +38,4 @@ videos/
 
 # testing
 .coverage*
-htmlcov/
+htmlcov/
diff --git a/README.md b/README.md
@@ -30,7 +30,8 @@ Please refer to our [documentation](https://haosulab.github.io/ManiSkill2) to le
 From pip:
 
 ```bash
-pip install mani-skill2
+# NOTE: This is a beta build at the moment. To install the beta build run
+pip install mani-skill2==0.6.0.dev0
 ```
 
 From github:

diff --git a/mani_skill2/agents/README.md b/mani_skill2/agents/README.md
@@ -0,0 +1,6 @@
+# Robots/Agents
+
+
+## How to create a new robot
+
+We recommend copying the files in `mani_skill2/robots/_template` and starting from there. 
diff --git a/mani_skill2/agents/base_agent.py b/mani_skill2/agents/base_agent.py
@@ -4,37 +4,29 @@
 from typing import Dict, Union
 
 import numpy as np
-import sapien.core as sapien
+import sapien
+import sapien.physx as physx
 from gymnasium import spaces
 
 from mani_skill2 import format_path
-from mani_skill2.sensors.camera import CameraConfig
-from mani_skill2.utils.sapien_utils import check_urdf_config, parse_urdf_config
+from mani_skill2.sensors.base_sensor import BaseSensor, BaseSensorConfig
+from mani_skill2.utils.sapien_utils import (
+    apply_urdf_config,
+    check_urdf_config,
+    parse_urdf_config,
+)
 
-from .base_controller import BaseController, CombinedController, ControllerConfig
-
-
-@dataclass
-class AgentConfig:
-    """Agent configuration.
-
-    Args:
-        urdf_path: path to URDF file. Support placeholders like {PACKAGE_ASSET_DIR}.
-        urdf_config: a dict to specify materials and simulation parameters when loading URDF from SAPIEN.
-        controllers: a dict of controller configurations
-        cameras: a dict of onboard camera configurations
-    """
-
-    urdf_path: str
-    urdf_config: dict
-    controllers: Dict[str, Union[ControllerConfig, Dict[str, ControllerConfig]]]
-    cameras: Dict[str, CameraConfig]
+from .controllers.base_controller import (
+    BaseController,
+    CombinedController,
+    ControllerConfig,
+)
 
 
 class BaseAgent:
     """Base class for agents.
 
-    Agent is an interface of the robot (sapien.Articulation).
+    Agent is an interface of an articulated robot (physx.PhysxArticulation).
 
     Args:
         scene (sapien.Scene): simulation scene instance.
@@ -44,45 +36,48 @@ class BaseAgent:
         config: agent configuration
     """
 
-    robot: sapien.Articulation
+    uid: str
+    robot: physx.PhysxArticulation
+
+    urdf_path: str
+    urdf_config: dict
+
+    controller_configs: Dict[str, Union[ControllerConfig, Dict[str, ControllerConfig]]]
     controllers: Dict[str, BaseController]
 
+    sensor_configs: Dict[str, BaseSensorConfig]
+    sensors: Dict[str, BaseSensor]
+
     def __init__(
         self,
         scene: sapien.Scene,
         control_freq: int,
         control_mode: str = None,
         fix_root_link=True,
-        config: AgentConfig = None,
     ):
         self.scene = scene
         self._control_freq = control_freq
 
-        self.config = config or self.get_default_config()
-
         # URDF
-        self.urdf_path = self.config.urdf_path
         self.fix_root_link = fix_root_link
-        self.urdf_config = self.config.urdf_config
 
         # Controller
-        self.controller_configs = self.config.controllers
         self.supported_control_modes = list(self.controller_configs.keys())
         if control_mode is None:
             control_mode = self.supported_control_modes[0]
         # The control mode after reset for consistency
         self._default_control_mode = control_mode
 
         self._load_articulation()
+        self._after_loading_articulation()
         self._setup_controllers()
         self.set_control_mode(control_mode)
         self._after_init()
 
-    @classmethod
-    def get_default_config(cls) -> AgentConfig:
-        raise NotImplementedError
-
     def _load_articulation(self):
+        """
+        Load the robot articulation
+        """
         loader = self.scene.create_urdf_loader()
         loader.fix_root_link = self.fix_root_link
 
@@ -92,29 +87,29 @@ def _load_articulation(self):
         check_urdf_config(urdf_config)
 
         # TODO(jigu): support loading multiple convex collision shapes
-        self.robot = loader.load(urdf_path, urdf_config)
+
+        apply_urdf_config(loader, urdf_config)
+        self.robot: physx.PhysxArticulation = loader.load(urdf_path)
         assert self.robot is not None, f"Fail to load URDF from {urdf_path}"
         self.robot.set_name(Path(urdf_path).stem)
 
         # Cache robot link ids
-        self.robot_link_ids = [link.get_id() for link in self.robot.get_links()]
+        self.robot_link_ids = [link.name for link in self.robot.get_links()]
 
-    def _setup_controllers(self):
-        self.controllers = OrderedDict()
-        for uid, config in self.controller_configs.items():
-            if isinstance(config, dict):
-                self.controllers[uid] = CombinedController(
-                    config, self.robot, self._control_freq
-                )
-            else:
-                self.controllers[uid] = config.controller_cls(
-                    config, self.robot, self._control_freq
-                )
+    def _after_loading_articulation(self):
+        """After loading articulation and before setting up controller. Not recommended, but is useful for when creating
+        robot classes that inherit controllers from another and only change which joints are controlled
+        """
+        pass
 
     def _after_init(self):
         """After initialization. E.g., caching the end-effector link."""
         pass
 
+    # -------------------------------------------------------------------------- #
+    # Controllers
+    # -------------------------------------------------------------------------- #
+
     @property
     def control_mode(self):
         """Get the currently activated controller uid."""
@@ -130,6 +125,21 @@ def set_control_mode(self, control_mode):
         self._control_mode = control_mode
         self.controller.reset()
 
+    def _setup_controllers(self):
+        """
+        Create and setup the controllers
+        """
+        self.controllers = OrderedDict()
+        for uid, config in self.controller_configs.items():
+            if isinstance(config, dict):
+                self.controllers[uid] = CombinedController(
+                    config, self.robot, self._control_freq
+                )
+            else:
+                self.controllers[uid] = config.controller_cls(
+                    config, self.robot, self._control_freq
+                )
+
     @property
     def controller(self):
         """Get currently activated controller."""
@@ -150,39 +160,38 @@ def action_space(self):
         else:
             return self.controller.action_space
 
-    def reset(self, init_qpos=None):
-        if init_qpos is not None:
-            self.robot.set_qpos(init_qpos)
-        self.robot.set_qvel(np.zeros(self.robot.dof))
-        self.robot.set_qacc(np.zeros(self.robot.dof))
-        self.robot.set_qf(np.zeros(self.robot.dof))
-        self.set_control_mode(self._default_control_mode)
-
     def set_action(self, action):
-        if np.isnan(action).any(): raise ValueError("Action cannot be NaN. Environment received:", action)
+        """
+        Set the agent's action which is to be executed in the next environment timestep
+        """
+        if np.isnan(action).any():
+            raise ValueError("Action cannot be NaN. Environment received:", action)
         self.controller.set_action(action)
 
     def before_simulation_step(self):
         self.controller.before_simulation_step()
 
     # -------------------------------------------------------------------------- #
-    # Observations
+    # Observations and State
     # -------------------------------------------------------------------------- #
     def get_proprioception(self):
+        """
+        Get the proprioceptive state of the agent.
+        """
         obs = OrderedDict(qpos=self.robot.get_qpos(), qvel=self.robot.get_qvel())
         controller_state = self.controller.get_state()
         if len(controller_state) > 0:
             obs.update(controller=controller_state)
         return obs
 
     def get_state(self) -> Dict:
-        """Get current state for MPC, including robot state and controller state"""
+        """Get current state, including robot state and controller state"""
         state = OrderedDict()
 
         # robot state
         root_link = self.robot.get_links()[0]
         state["robot_root_pose"] = root_link.get_pose()
-        state["robot_root_vel"] = root_link.get_velocity()
+        state["robot_root_vel"] = root_link.get_linear_velocity()
         state["robot_root_qvel"] = root_link.get_angular_velocity()
         state["robot_qpos"] = self.robot.get_qpos()
         state["robot_qvel"] = self.robot.get_qvel()
@@ -204,3 +213,34 @@ def set_state(self, state: Dict, ignore_controller=False):
 
         if not ignore_controller and "controller" in state:
             self.controller.set_state(state["controller"])
+
+    # -------------------------------------------------------------------------- #
+    # Other
+    # -------------------------------------------------------------------------- #
+    def reset(self, init_qpos=None):
+        """
+        Reset the robot to a rest position or a given q-position
+        """
+        if init_qpos is not None:
+            self.robot.set_qpos(init_qpos)
+        self.robot.set_qvel(np.zeros(self.robot.dof))
+        self.robot.set_qacc(np.zeros(self.robot.dof))
+        self.robot.set_qf(np.zeros(self.robot.dof))
+        self.set_control_mode(self._default_control_mode)
+
+    # -------------------------------------------------------------------------- #
+    # Optional per-agent APIs, implemented depending on agent affordances
+    # -------------------------------------------------------------------------- #
+    def is_grasping(self, object: Union[sapien.Entity, None] = None):
+        """
+        Check if this agent is grasping an object or grasping anything at all
+
+        Args:
+            object (sapien.Entity | None):
+                If object is a sapien.Entity, this function checks grasping against that. If it is none, the function checks if the
+                agent is grasping anything at all.
+
+        Returns:
+            True if agent is grasping object. False otherwise. If object is None, returns True if agent is grasping something, False if agent is grasping nothing.
+        """
+        raise NotImplementedError()
diff --git a/mani_skill2/agents/configs/__init__.py b/mani_skill2/agents/configs/__init__.py
diff --git a/mani_skill2/agents/configs/mobile_panda/__init__.py b/mani_skill2/agents/configs/mobile_panda/__init__.py
diff --git a/mani_skill2/agents/configs/panda/__init__.py b/mani_skill2/agents/configs/panda/__init__.py
diff --git a/mani_skill2/agents/configs/xmate3/__init__.py b/mani_skill2/agents/configs/xmate3/__init__.py