From 13e97bec741eb10fdfcac1eedc8fd451959372b0 Mon Sep 17 00:00:00 2001 From: Kaixing Xiao Date: Fri, 16 Aug 2024 04:50:14 +0200 Subject: [PATCH] Add doc string --- deformable_gym/envs/mujoco/asset_manager.py | 55 ++++++++++++++++ deformable_gym/envs/mujoco/base_mjenv.py | 45 ++++++++++++- deformable_gym/envs/mujoco/grasp_env.py | 71 +++++++++++++++++---- deformable_gym/objects/mj_object.py | 55 +++++++++++++--- deformable_gym/robots/mj_robot.py | 67 +++++++++++++++++++ 5 files changed, 267 insertions(+), 26 deletions(-) diff --git a/deformable_gym/envs/mujoco/asset_manager.py b/deformable_gym/envs/mujoco/asset_manager.py index 0914550..83a37f5 100644 --- a/deformable_gym/envs/mujoco/asset_manager.py +++ b/deformable_gym/envs/mujoco/asset_manager.py @@ -14,6 +14,11 @@ class AssetManager: + """ + The AssetManager class manages the loading and combining of MuJoCo models (robots and objects) + into a single simulation scene. It provides methods to load individual robot or object models, + create complex scenes, and save the resulting XML configuration. + """ def __init__(self) -> None: self.assets_dir = os.path.join(os.path.dirname(__file__), "assets") @@ -22,6 +27,19 @@ def __init__(self) -> None: self.objects = OBJECTS def load_asset(self, name: str) -> mujoco.MjModel: + """ + Loads a asset (robot or object mujoco Model) based on its name. + + Args: + name (str): The name of the robot or object to be loaded. Must be a key in either + self.robots or self.objects. + + Returns: + mujoco.MjModel: The loaded MuJoCo model. + + Raises: + AssertionError: If the specified name is not found in either self.robots or self.objects. + """ assert ( name in self.robots or name in self.objects ), f"Model {name} not found.\n available: {list(self.robots.keys()) + list(self.objects.keys())}" @@ -31,6 +49,20 @@ def load_asset(self, name: str) -> mujoco.MjModel: return model def create_scene(self, robot_name: str, obj_name: str) -> str: + """ + Creates an MJCF string representing a MuJoCo scene that includes a robot and an object. + + Args: + robot_name (str): The name of the robot to include in the scene. Must be a key in self.robots. + obj_name (str): The name of the object to include in the scene. Must be a key in self.objects. + + Returns: + str: A MJCF string representing the combined MuJoCo scene. + + Raises: + AssertionError: If the specified robot_name is not found in self.robots. + AssertionError: If the specified obj_name is not found in self.objects. + """ assert ( robot_name in self.robots ), f"Robot {robot_name} not found.\n available: {list(self.robots.keys())}" @@ -63,6 +95,15 @@ def create_scene(self, robot_name: str, obj_name: str) -> str: # return scene def _get_full_path(self, file: str) -> str: + """ + Generates the full path to a file within the assets directory. + + Args: + file (str): The file name (with extension) for which to generate the full path. + + Returns: + str: The full path to the specified file within the assets directory. + """ return os.path.join(self.assets_dir, file) @staticmethod @@ -72,6 +113,20 @@ def include_mjcf( *, meshdir: Optional[str] = None, ) -> str: + """ + Generates an XML string for a MuJoCo scene by including additional MJCF files within a base MJCF file. + + Args: + base_path (str): The file path to the base MJCF file. + include_path (Union[str, Sequence[str]]): A string or list of strings representing file paths + to MJCF files to be included in the base file. + meshdir (Optional[str]): A string representing the path to the directory containing mesh files. + If provided, this path is added to the meshdir attribute of the compiler + element in the MJCF XML. + + Returns: + str: An XML string representing the combined MJCF file. + """ tree = ET.parse(base_path) root = tree.getroot() if isinstance(include_path, list) or isinstance(include_path, tuple): diff --git a/deformable_gym/envs/mujoco/base_mjenv.py b/deformable_gym/envs/mujoco/base_mjenv.py index a8d8a20..730b2cd 100644 --- a/deformable_gym/envs/mujoco/base_mjenv.py +++ b/deformable_gym/envs/mujoco/base_mjenv.py @@ -39,16 +39,41 @@ def __init__( self.action_space = self._get_action_space() def _create_scene(self, robot_name: str, obj_name: str) -> str: - manager = AssetManager() - return manager.create_scene(robot_name, obj_name) + """ + Creates the simulation scene by combining the robot and object models. + + Args: + robot_name (str): The name of the robot to include in the scene. + obj_name (str): The name of the object to include in the scene. + + Returns: + str: The MJCF XML string representing the combined simulation scene. + """ + + return AssetManager().create_scene(robot_name, obj_name) def _get_action_space(self) -> spaces.Box: + """ + Defines the action space for the environment based on the robot's control range. + + Returns: + spaces.Box: A continuous space representing the possible actions the agent can take. + """ + n_act = self.robot.nact low = self.robot.ctrl_range[:, 0].copy() high = self.robot.ctrl_range[:, 1].copy() return spaces.Box(low=low, high=high, shape=(n_act,), dtype=np.float64) def _get_observation_space(self) -> spaces.Box: + """ + Defines the observation space for the environment, including the robot's joint positions + and, optionally, the object's position. + + Returns: + spaces.Box: A continuous space representing the state observations available to the agent. + """ + nq = self.robot.nq low = -np.inf # TODO: joint space range high = np.inf @@ -65,8 +90,16 @@ def reset( options: Optional[dict] = None, ) -> tuple[NDArray[np.float64], dict[str, Any]]: """ - Reset the environment to the initial state. + Resets the environment to its initial state. + + Args: + seed (Optional[int], optional): A random seed for resetting the environment. Default is None. + options (Optional[dict], optional): Additional options for resetting the environment. Default is None. + + Returns: + tuple: A tuple containing the initial observation and an empty info dictionary. """ + super().reset(seed=seed, options=options) self.model, _ = mju.load_model_from_string(self.scene) mujoco.mj_resetData(self.model, self.data) @@ -93,6 +126,12 @@ def _set_state( mujoco.mj_forward(self.model, self.data) def _load_keyframe(self, frame_name: str): + """ + Loads a predefined keyframe and sets the environment's state to it. + + Args: + frame_name (str): The name of the keyframe to load. + """ frame = self.model.keyframe(frame_name) qpos = frame.qpos.copy() qvel = frame.qvel.copy() diff --git a/deformable_gym/envs/mujoco/grasp_env.py b/deformable_gym/envs/mujoco/grasp_env.py index 6d04e93..8884b4f 100644 --- a/deformable_gym/envs/mujoco/grasp_env.py +++ b/deformable_gym/envs/mujoco/grasp_env.py @@ -10,6 +10,9 @@ class GraspEnv(BaseMJEnv): + """ + A custom MuJoCo environment for a grasping task, where a robot attempts to grasp an object. + """ def __init__( self, @@ -35,14 +38,18 @@ def __init__( def reset(self, *, seed=None, options=None) -> Tuple[NDArray, Dict]: super().reset(seed=seed, options=options) - self.robot.set_pose( - self.model, self.data, self.robot.init_pose[self.object.name] - ) observation = self._get_observation() info = self._get_info() return observation, info def _get_observation(self) -> NDArray: + """ + The observation includes the robot's joint qpos in their generalized coordinates, + and optionally, the object's position. + + Returns: + NDArray: A numpy array containing the observation. + """ robot_qpos = self.robot.get_qpos(self.model, self.data) if self.observable_object_pos: obj_pos = self.object.get_current_com(self.data) @@ -52,13 +59,19 @@ def _get_observation(self) -> NDArray: return obs def _pause_simulation(self, time: float) -> None: - """Step Mujoco Simulation for a given time. + """ + Steps the MuJoCo simulation for a specified amount of time. + + certain constraints of object will be disabled and all joints will be freezed + before stepping the simulation to allow the robot to grasp the object without + external influences. Args: - time (float): simulation time in seconds + time (float): The duration in seconds for which to step the simulation. """ - # mju.remove_geom(self.model, self.data, "platform") # might be off here... - self.object.disable_eq_constraint(self.model, self.data) + mju.disable_equality_constraint( + self.model, self.data, *self.object.eq_constraints_to_disable + ) mju.disable_joint(self.model, self.data, *self.robot.joints) start_time = self.data.time while self.data.time - start_time < time: @@ -67,17 +80,22 @@ def _pause_simulation(self, time: float) -> None: self.render() def _get_reward(self, terminated: bool) -> int: - """Calculate reward by removing the platform and check if object falls to the ground. - 0 reward: max_sim_time is not reached yet - -1 reward: object falls to the ground after removing the platform - 1 reward: object is grasped successfully by the robot hand + """ + Calculates the reward based on the robot's success in grasping the object. + + The reward is calculated after removing all fixed constraints and checking + if the object remains grasped by the robot. Args: - terminated (bool): if episode is terminated + terminated (bool): Whether the episode has terminated. Returns: - int: reward gotten per step + int: The reward for the current step. Possible values are: + - 0: The episode has not yet terminated. + - 1: The object is successfully grasped by the robot. + - -1: The object falls to the ground. """ + if not terminated: return 0 self._pause_simulation(1) @@ -88,17 +106,42 @@ def _get_reward(self, terminated: bool) -> int: return -1 def _is_terminated(self, sim_time: float) -> bool: + """ + Determines whether the episode has terminated based on the simulation time. + + Args: + sim_time (float): The current simulation time. + + Returns: + bool: True if the simulation time has exceeded the maximum allowed time, otherwise False. + """ return sim_time >= self.max_sim_time def _is_truncated(self) -> bool: return False def _get_info(self) -> Dict: - if self.viewer is not None: + """ + If the GUI viewer is running, this method will return a dictionary indicating that. + + Returns: + Dict: A dictionary containing information about the environment. + """ + if self.gui and self.viewer is not None: return {"is_viewer_running": self.viewer.is_running()} return {} def step(self, ctrl: ArrayLike) -> Tuple[NDArray, int, bool, bool, Dict]: + """ + Advances the simulation applying the given control input to the robot + + Args: + ctrl (ArrayLike): The control input to be applied to the robot. + + Returns: + Tuple[NDArray, int, bool, bool, Dict]: observation, reward, termination flag, + truncation flag, and an info. + """ sim_time = self.data.time self.robot.set_ctrl(self.model, self.data, ctrl) mujoco.mj_step(self.model, self.data) diff --git a/deformable_gym/objects/mj_object.py b/deformable_gym/objects/mj_object.py index bc1ee9c..ce3a17c 100644 --- a/deformable_gym/objects/mj_object.py +++ b/deformable_gym/objects/mj_object.py @@ -9,6 +9,20 @@ class MJObject: + """ + A class representing a physical object in a MuJoCo simulation. + It provides methods to set the object's pose, get the center of mass, and manage equality constraints. + + Attributes: + name (str): The name of the object, corresponding to its definition in the XML model. + model (mujoco.MjModel): The MuJoCo model object representing the object in the simulation. + eq_constraints (List[str]): A list of equality constraints associated with this object. + These constraints are typically used to enforce specific relationships between bodies, + such as keeping them at a fixed distance or maintaining an orientation. + eq_constraints_to_disable (List[Union[str, None]]): A list of equality constraints that are marked to be disabled for this object + in a specific occasion such as making the object free to move in the simulation for floating objects. + This list can be customized by subclasses to define which constraints should be ignored or temporarily disabled during the simulation. + """ def __init__(self, name: str) -> None: self.name = name @@ -23,8 +37,18 @@ def eq_constraints_to_disable(self) -> List[Union[str, None]]: return [] def _load_model(self, name: str) -> mujoco.MjModel: - manager = AssetManager() - return manager.load_asset(name) + """ + Load the MuJoCo model for the object. + + This method uses the `AssetManager` to load the XML model file corresponding to the object. + + Args: + name (str): The name of the object to load. + + Returns: + mujoco.MjModel: The loaded MuJoCo model for the object. + """ + return AssetManager().load_asset(name) def set_pose( self, @@ -32,19 +56,32 @@ def set_pose( data: mujoco.MjData, pose: Pose, ) -> None: + """ + Set the pose (position and orientation) of the object in the simulation. + + This method updates the position and orientation of the object's body in the simulation + and recalculates the simulation state. + + Args: + model (mujoco.MjModel): The MuJoCo model object containing the object's configuration. + data (mujoco.MjData): The MuJoCo data object containing the current simulation state. + pose (Pose): A Pose object containing the desired position and orientation for the object. + """ model.body(self.name).pos[:] = pose.position model.body(self.name).quat[:] = pose.orientation mujoco.mj_forward(model, data) def get_current_com(self, data: mujoco.MjData) -> NDArray: - return data.body(self.name).ipos + """ + Get the current center of mass (COM) position of the object. - def disable_eq_constraint( - self, model: mujoco.MjModel, data: mujoco.MjData, *name: str - ) -> None: - if len(name) == 0: - name = self.eq_constraints_to_disable - mju.disable_equality_constraint(model, data, *name) + Args: + data (mujoco.MjData): The MuJoCo data object containing the current simulation state. + + Returns: + NDArray: An array representing the position of the object's center of mass. + """ + return data.body(self.name).ipos class InsoleFixed(MJObject): diff --git a/deformable_gym/robots/mj_robot.py b/deformable_gym/robots/mj_robot.py index 750a7e7..eca13a9 100644 --- a/deformable_gym/robots/mj_robot.py +++ b/deformable_gym/robots/mj_robot.py @@ -11,6 +11,24 @@ class MJRobot(ABC): + """ + Abstract base class for creating robot instances in a MuJoCo simulation environment. + + This class provides a set of methods to manage robot configurations, including + initialization, accessing joint and actuator information, and setting robot poses and controls. + + Attributes: + init_pos (dict): A dictionary that can store initial positions for different robot models. + This can be overridden by subclasses to provide specific initial positions. + name (str): The name of the robot model. + model (mujoco.MjModel): The MuJoCo model of the robot, loaded from the corresponding XML file. + nq (int): Number of generalized coordinates (joints) in the robot model. + dof (int): Degrees of freedom in the robot model. + nact (int): Number of actuators in the robot model. + ctrl_range (NDArray): The control range for each actuator, indicating the min and max values. + joints (List[str]): A list of joint names for the robot. + actuators (List[str]): A list of actuator names for the robot. + """ init_pos = {} @@ -32,9 +50,31 @@ def actuators(self) -> List[str]: return mju.get_actuator_names(self.model) def get_qpos(self, model: mujoco.MjModel, data: mujoco.MjData) -> NDArray: + """ + Get the current joint positions (qpos) for the robot. + + Args: + model (mujoco.MjModel): The MuJoCo model object containing the robot's configuration. + data (mujoco.MjData): The MuJoCo data object containing the current simulation state. + + Returns: + NDArray: An array of joint positions for the robot. + """ + return mju.get_joint_qpos(model, data, *self.joints) def get_qvel(self, model: mujoco.MjModel, data: mujoco.MjData) -> NDArray: + """ + Get the current joint velocities (qvel) for the robot. + + Args: + model (mujoco.MjModel): The MuJoCo model object containing the robot's configuration. + data (mujoco.MjData): The MuJoCo data object containing the current simulation state. + + Returns: + NDArray: An array of joint velocities for the robot. + """ + return mju.get_joint_qvel(model, data, *self.joints) def set_pose( @@ -43,6 +83,18 @@ def set_pose( data: mujoco.MjData, pose: Pose, ) -> None: + """ + Set the pose (position and orientation) of the robot's base. + + This method sets the position and orientation of the robot's base body in the simulation, + and then recalculates the simulation state using `mujoco.mj_forward`. + + Args: + model (mujoco.MjModel): The MuJoCo model object containing the robot's configuration. + data (mujoco.MjData): The MuJoCo data object containing the current simulation state. + pose (Pose): A Pose object containing the desired position and orientation for the robot's base. + """ + model.body(self.name).pos[:] = pose.position model.body(self.name).quat[:] = pose.orientation mujoco.mj_forward(model, data) @@ -50,6 +102,21 @@ def set_pose( def set_ctrl( self, model: mujoco.MjModel, data: mujoco.MjData, ctrl: ArrayLike ) -> None: + """ + Apply control inputs to the robot's actuators. + + This method sets the control input for each actuator based on the provided control vector. + The control vector must have a length equal to the number of actuators. + + Args: + model (mujoco.MjModel): The MuJoCo model object containing the robot's configuration. + data (mujoco.MjData): The MuJoCo data object containing the current simulation state. + ctrl (ArrayLike): An array of control inputs for the actuators. + + Raises: + AssertionError: If the length of the control vector does not match the number of actuators. + """ + assert ( len(ctrl) == self.nact ), f"Control vector should have length {self.nact}"