copy code over

.gitignore

@@ -38,4 +38,4 @@ videos/
 
 # testing
 .coverage*
-htmlcov/
+htmlcov/

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:

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. 

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,29 +36,32 @@ 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]
@@ -78,11 +73,10 @@ def __init__(
         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 +86,23 @@ 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()]
-
-    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
-                )
+        self.robot_link_ids = [link.name for link in self.robot.get_links()]
 
     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 +118,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 +153,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 +206,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()

mani_skill2/agents/base_controller.py → mani_skill2/agents/controllers/base_controller.py

@@ -3,27 +3,28 @@
 from typing import Dict, List
 
 import numpy as np
-import sapien.core as sapien
+import sapien
+import sapien.physx as physx
 from gymnasium import spaces
 
 from mani_skill2.utils.common import clip_and_scale_action, normalize_action_space
 
-from .utils import flatten_action_spaces, get_active_joint_indices, get_active_joints
+from mani_skill2.agents.utils import flatten_action_spaces, get_active_joint_indices, get_active_joints
 
 
 class BaseController:
     """Base class for controllers.
     The controller is an interface for the robot to interact with the environment.
     """
 
-    joints: List[sapien.Joint]  # active joints controlled
+    joints: List[physx.PhysxArticulationJoint]  # active joints controlled
     joint_indices: List[int]  # indices of active joints controlled
     action_space: spaces.Space
 
     def __init__(
         self,
         config: "ControllerConfig",
-        articulation: sapien.Articulation,
+        articulation: physx.PhysxArticulation,
         control_freq: int,
         sim_freq: int = None,
     ):
@@ -33,8 +34,7 @@ def __init__(
 
         # For action interpolation
         if sim_freq is None:  # infer from scene
-            # TODO(jigu): update sapien interface to avoid this workaround
-            sim_timestep = self.articulation.get_builder().get_scene().get_timestep()
+            sim_timestep = self.articulation.root.entity.get_scene().timestep
             sim_freq = round(1.0 / sim_timestep)
         # Number of simulation steps per control step
         self._sim_steps = sim_freq // control_freq
@@ -143,7 +143,7 @@ class DictController(BaseController):
     def __init__(
         self,
         configs: Dict[str, ControllerConfig],
-        articulation: sapien.Articulation,
+        articulation: physx.PhysxArticulation,
         control_freq: int,
         sim_freq: int = None,
         balance_passive_force=True,
@@ -203,7 +203,9 @@ def set_action(self, action: Dict[str, np.ndarray]):
 
     def before_simulation_step(self):
         if self.balance_passive_force:
-            qf = self.articulation.compute_passive_force(external=False)
+            qf = self.articulation.compute_passive_force(
+                gravity=True, coriolis_and_centrifugal=True
+            )
         else:
             qf = np.zeros(self.articulation.dof)
         for controller in self.controllers.values():

mani_skill2/agents/controllers/passive_controller.py

@@ -4,10 +4,14 @@
 import numpy as np
 from gymnasium import spaces
 
-from ..base_controller import BaseController, ControllerConfig
+from .base_controller import BaseController, ControllerConfig
 
 
 class PassiveController(BaseController):
+    """
+    Passive controller that does not do anything
+    """
+
     config: "PassiveControllerConfig"
 
     def set_drive_property(self):

mani_skill2/agents/controllers/pd_ee_pose.py

@@ -2,14 +2,14 @@
 from typing import Sequence, Union
 
 import numpy as np
-import sapien.core as sapien
+import sapien
 from gymnasium import spaces
 from scipy.spatial.transform import Rotation
 
 from mani_skill2.utils.common import clip_and_scale_action
-from mani_skill2.utils.sapien_utils import get_entity_by_name, vectorize_pose
+from mani_skill2.utils.sapien_utils import get_obj_by_name, vectorize_pose
 
-from ..base_controller import BaseController, ControllerConfig
+from .base_controller import BaseController, ControllerConfig
 from .pd_joint_pos import PDJointPosController
 
 
@@ -26,7 +26,7 @@ def _initialize_joints(self):
         self.qmask[self.joint_indices] = 1
 
         if self.config.ee_link:
-            self.ee_link = get_entity_by_name(
+            self.ee_link = get_obj_by_name(
                 self.articulation.get_links(), self.config.ee_link
             )
         else:
@@ -50,7 +50,7 @@ def ee_pose(self):
     @property
     def ee_pose_at_base(self):
         to_base = self.articulation.pose.inv()
-        return to_base.transform(self.ee_pose)
+        return to_base * (self.ee_pose)
 
     def reset(self):
         super().reset()