From 8273a1f9a6a40638b3404c32191b465feb443af5 Mon Sep 17 00:00:00 2001 From: Filippo Luca Ferretti Date: Wed, 27 Mar 2024 21:54:19 +0100 Subject: [PATCH 01/21] Remove old modules --- docs/modules/high_level.rst | 30 ---------------------------- docs/modules/physics.rst | 40 ------------------------------------- docs/modules/simulation.rst | 34 ------------------------------- 3 files changed, 104 deletions(-) delete mode 100644 docs/modules/high_level.rst delete mode 100644 docs/modules/physics.rst delete mode 100644 docs/modules/simulation.rst diff --git a/docs/modules/high_level.rst b/docs/modules/high_level.rst deleted file mode 100644 index 52528d708..000000000 --- a/docs/modules/high_level.rst +++ /dev/null @@ -1,30 +0,0 @@ -High Level -========== - -.. currentmodule:: jaxsim.high_level - - -Model -~~~~~ - -.. automodule:: jaxsim.high_level.model - :members: - -Joint -~~~~~ - -.. automodule:: jaxsim.high_level.joint - :members: - -Link -~~~~~ -.. automodule:: jaxsim.high_level.link - :members: - -Common -~~~~~~ - -.. autoflag:: jaxsim.high_level.common.VelRepr - :members: - - diff --git a/docs/modules/physics.rst b/docs/modules/physics.rst deleted file mode 100644 index 1050f02ce..000000000 --- a/docs/modules/physics.rst +++ /dev/null @@ -1,40 +0,0 @@ -Physics -======= - -Algos ------ - -.. autofunction:: jaxsim.physics.algos.aba.aba - -.. autofunction:: jaxsim.physics.algos.crba.crba - -.. autofunction:: jaxsim.physics.algos.forward_kinematics.forward_kinematics - -.. autofunction:: jaxsim.physics.algos.jacobian.jacobian - -.. autofunction:: jaxsim.physics.algos.rnea.rnea - -.. automodule:: jaxsim.physics.algos.soft_contacts - :members: - -.. automodule:: jaxsim.physics.algos.terrain - :members: - -.. automodule:: jaxsim.physics.algos.utils - :members: - -Model ------ - -.. automodule:: jaxsim.physics.model - :members: - -.. automodule:: jaxsim.physics.model.ground_contact - :members: - -.. automodule:: jaxsim.physics.model.physics_model_state - :members: - -.. automodule:: jaxsim.physics.model.physics_model - :members: - diff --git a/docs/modules/simulation.rst b/docs/modules/simulation.rst deleted file mode 100644 index 0de03f483..000000000 --- a/docs/modules/simulation.rst +++ /dev/null @@ -1,34 +0,0 @@ -Simulation -========== - -.. automodule:: jaxsim.simulation - -Simulator -~~~~~~~~~ - -.. autoclass:: jaxsim.simulation.simulator.SimulatorData - :members: - -.. autoclass:: jaxsim.simulation.simulator.JaxSim - :members: - -Simulator Callbacks -~~~~~~~~~~~~~~~~~~~ - -.. currentmodule:: jaxsim.simulation.simulator_callbacks - -.. autosummary:: - CallbackHandler - ConfigureCallback - PostStepCallback - PreStepCallback - SimulatorCallback - -ODE -~~~ - -.. automodule:: jaxsim.simulation.ode - :members: - -.. autoflag:: jaxsim.simulation.ode_integration.IntegratorType - :members: \ No newline at end of file From 1dd142f8b5357b8c2308d1f0e448643f24aa7e05 Mon Sep 17 00:00:00 2001 From: Filippo Luca Ferretti Date: Wed, 27 Mar 2024 21:54:36 +0100 Subject: [PATCH 02/21] Add `api` module --- docs/modules/api.rst | 69 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 69 insertions(+) create mode 100644 docs/modules/api.rst diff --git a/docs/modules/api.rst b/docs/modules/api.rst new file mode 100644 index 000000000..3f79d76f4 --- /dev/null +++ b/docs/modules/api.rst @@ -0,0 +1,69 @@ +API +========== + +.. currentmodule:: jaxsim.api + + +Model +~~~~~ + +.. automodule:: jaxsim.api.model + :members: + +Data +~~~~ + +.. automodule:: jaxsim.api.data + :members: + +Contact +~~~~~~~ + +.. automodule:: jaxsim.api.contact + :members: + +KinDynParameters +~~~~~~~~~~~~~~~~ + +.. automodule:: jaxsim.api.kin_dyn_parameters + :members: + +Joint +~~~~~ + +.. automodule:: jaxsim.api.joint + :members: + +Link +~~~~~ +.. automodule:: jaxsim.api.link + :members: + +CoM +~~~ +.. automodule:: jaxsim.api.com + :members: + +ODE Data +~~~~~~~~ + +.. automodule:: jaxsim.api.ode_data + :members: + +.. automodule:: jaxsim.api.ode + :members: + +References +~~~~~~~~~~ + +.. automodule:: jaxsim.api.references + :members: + +Common +~~~~~~ + +.. autoflag:: jaxsim.api.common.VelRepr + :members: + +.. autoclass:: jaxsim.api.common.ModelDataWithVelocityRepresentation + :members: From 4bc9ec4e6a902757bafc51276998c3e4e2781547 Mon Sep 17 00:00:00 2001 From: Filippo Luca Ferretti Date: Wed, 27 Mar 2024 21:55:05 +0100 Subject: [PATCH 03/21] Add `integrators` module --- docs/modules/integrators.rst | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) create mode 100644 docs/modules/integrators.rst diff --git a/docs/modules/integrators.rst b/docs/modules/integrators.rst new file mode 100644 index 000000000..93b4ff555 --- /dev/null +++ b/docs/modules/integrators.rst @@ -0,0 +1,19 @@ +Integrators +=========== + +.. currentmodule:: jaxsim.integrators + +Common +~~~~~~ + +.. automodule:: jaxsim.integrators.common + +Fixed Step +~~~~~~~~~~ + +.. automodule:: jaxsim.integrators.fixed_step + +Variable Step +~~~~~~~~~~~~~ + +.. automodule:: jaxsim.integrators.variable_step From 832a2d6ff6245c9958ac56651f118f7a51295b15 Mon Sep 17 00:00:00 2001 From: Filippo Luca Ferretti Date: Wed, 27 Mar 2024 21:55:30 +0100 Subject: [PATCH 04/21] Add `rbda` module --- docs/modules/rbda.rst | 41 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) create mode 100644 docs/modules/rbda.rst diff --git a/docs/modules/rbda.rst b/docs/modules/rbda.rst new file mode 100644 index 000000000..0ac63e7ce --- /dev/null +++ b/docs/modules/rbda.rst @@ -0,0 +1,41 @@ +Rigid Body Dynamics Algorithms +============================== + +This module provides a set of algorithms for rigid body dynamics. + +.. currentmodule:: jaxsim.rbda + + +Articulated Body Algorithm +~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. automodule:: jaxsim.rbda.aba + +Collision Detection +~~~~~~~~~~~~~~~~~~~ + +.. automodule:: jaxsim.rbda.collidable_points + +Composite Rigid Body Algorithm +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. automodule:: jaxsim.rbda.crba + +Forward Kinetmatics +~~~~~~~~~~~~~~~~~~~ + +.. automodule:: jaxsim.rbda.forward_kinematics + +Jacobians +~~~~~~~~~ + +.. automodule:: jaxsim.rbda.jacobian + +Soft Contacts +~~~~~~~~~~~~~ + +.. automodule:: jaxsim.rbda.soft_contacts + +Utilities + +.. automodule:: jaxsim.rbda.utils From 785858e9baf23396307a65a36c37b254126dff6f Mon Sep 17 00:00:00 2001 From: Filippo Luca Ferretti Date: Wed, 27 Mar 2024 21:58:17 +0100 Subject: [PATCH 05/21] Add `mujoco` module --- docs/modules/mujoco.rst | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) create mode 100644 docs/modules/mujoco.rst diff --git a/docs/modules/mujoco.rst b/docs/modules/mujoco.rst new file mode 100644 index 000000000..a9266cf2b --- /dev/null +++ b/docs/modules/mujoco.rst @@ -0,0 +1,27 @@ +MuJoCo Visualizer +================== + +JAXsim provides a simple interface with MuJoCo's visualizer. The visualizer is +a separate process that communicates with the main simulation process. This +allows for the simulation to run at full speed while the visualizer can run at +a different frame rate. + +.. currentmodule:: jaxsim.mujoco.visualizer + +Loaders +~~~~~~~ + +.. automodule:: jaxsim.mujoco.loaders + :members: + +Model +~~~~~ + +.. automodule:: jaxsim.mujoco.model + :members: + +Visualizer +~~~~~~~~~~ + +.. automodule:: jaxsim.mujoco.visualizer + :members: From 020c6826c4dbdd512dc155575b32a1be8c80eb41 Mon Sep 17 00:00:00 2001 From: Filippo Luca Ferretti Date: Wed, 27 Mar 2024 21:58:46 +0100 Subject: [PATCH 06/21] Update `utils` module --- docs/modules/utils.rst | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/docs/modules/utils.rst b/docs/modules/utils.rst index 49470257e..93a8fc5d9 100644 --- a/docs/modules/utils.rst +++ b/docs/modules/utils.rst @@ -1,7 +1,4 @@ -.. _utils: - - -Utils +Utils ===== .. automodule:: jaxsim.utils From ffc0ac1758ab0c268581bab399bb5a5a40d8ada9 Mon Sep 17 00:00:00 2001 From: Filippo Luca Ferretti Date: Wed, 27 Mar 2024 21:59:14 +0100 Subject: [PATCH 07/21] Update theme --- docs/conf.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/conf.py b/docs/conf.py index 7ba001bd5..7a4eb3c30 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -64,7 +64,7 @@ def _recursive_add_annotations_import(): language = "en" -html_theme = "sphinx_rtd_theme" +html_theme = "sphinx_book_theme" templates_path = ["_templates"] From 749a3c09af880cb89304cc4c498c7f94f477adc2 Mon Sep 17 00:00:00 2001 From: Filippo Luca Ferretti Date: Wed, 27 Mar 2024 21:59:40 +0100 Subject: [PATCH 08/21] Update index --- docs/index.rst | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/docs/index.rst b/docs/index.rst index 73ad14402..52c946397 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -40,12 +40,12 @@ Features :maxdepth: 2 :caption: JAXsim API - modules/high_level + modules/api + modules/integrators + modules/rbda modules/math + modules/mujoco modules/parsers - modules/physics - modules/simulation - modules/typing modules/utils @@ -76,7 +76,7 @@ Here below we summarize the differences between the projects: Contributing ------------ -Pull requests are welcome. +Pull requests are welcome. For major changes, please open an issue first to discuss what you would like to change. Citing From 67c5cde2999b2f062f3c62d883c4df53684203c4 Mon Sep 17 00:00:00 2001 From: Filippo Luca Ferretti Date: Wed, 27 Mar 2024 22:14:42 +0100 Subject: [PATCH 09/21] Update docs environment --- environment.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/environment.yml b/environment.yml index 5f8ce23ef..e182c4c62 100644 --- a/environment.yml +++ b/environment.yml @@ -42,4 +42,5 @@ dependencies: - sphinx-jinja2-compat - sphinx-multiversion - sphinx_rtd_theme + - sphinx-book-theme - sphinx-toolbox From 88f9e770aef4a0265da929d472ae56fd5a377987 Mon Sep 17 00:00:00 2001 From: Filippo Luca Ferretti Date: Wed, 27 Mar 2024 23:56:29 +0100 Subject: [PATCH 10/21] Add missing members --- docs/modules/integrators.rst | 3 +++ docs/modules/mujoco.rst | 2 +- docs/modules/rbda.rst | 8 ++++++++ 3 files changed, 12 insertions(+), 1 deletion(-) diff --git a/docs/modules/integrators.rst b/docs/modules/integrators.rst index 93b4ff555..2a1e5327d 100644 --- a/docs/modules/integrators.rst +++ b/docs/modules/integrators.rst @@ -7,13 +7,16 @@ Common ~~~~~~ .. automodule:: jaxsim.integrators.common + :members: Fixed Step ~~~~~~~~~~ .. automodule:: jaxsim.integrators.fixed_step + :members: Variable Step ~~~~~~~~~~~~~ .. automodule:: jaxsim.integrators.variable_step + :members: diff --git a/docs/modules/mujoco.rst b/docs/modules/mujoco.rst index a9266cf2b..8b304837e 100644 --- a/docs/modules/mujoco.rst +++ b/docs/modules/mujoco.rst @@ -6,7 +6,7 @@ a separate process that communicates with the main simulation process. This allows for the simulation to run at full speed while the visualizer can run at a different frame rate. -.. currentmodule:: jaxsim.mujoco.visualizer +.. currentmodule:: jaxsim.mujoco Loaders ~~~~~~~ diff --git a/docs/modules/rbda.rst b/docs/modules/rbda.rst index 0ac63e7ce..60e69be97 100644 --- a/docs/modules/rbda.rst +++ b/docs/modules/rbda.rst @@ -10,32 +10,40 @@ Articulated Body Algorithm ~~~~~~~~~~~~~~~~~~~~~~~~~~ .. automodule:: jaxsim.rbda.aba + :members: Collision Detection ~~~~~~~~~~~~~~~~~~~ .. automodule:: jaxsim.rbda.collidable_points + :members: Composite Rigid Body Algorithm ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ .. automodule:: jaxsim.rbda.crba + :members: Forward Kinetmatics ~~~~~~~~~~~~~~~~~~~ .. automodule:: jaxsim.rbda.forward_kinematics + :members: Jacobians ~~~~~~~~~ .. automodule:: jaxsim.rbda.jacobian + :members: Soft Contacts ~~~~~~~~~~~~~ .. automodule:: jaxsim.rbda.soft_contacts + :members: Utilities +~~~~~~~~~ .. automodule:: jaxsim.rbda.utils + :members: From c3f57134c6f347597aa4811524a67cd34491a03f Mon Sep 17 00:00:00 2001 From: Filippo Luca Ferretti Date: Thu, 28 Mar 2024 11:27:00 +0100 Subject: [PATCH 11/21] Add docstrings to mujoco module --- src/jaxsim/mujoco/loaders.py | 69 ++++++++++++++++++++++++++++++--- src/jaxsim/mujoco/model.py | 57 ++++++++++++++++++--------- src/jaxsim/mujoco/visualizer.py | 34 +++++++++++----- 3 files changed, 126 insertions(+), 34 deletions(-) diff --git a/src/jaxsim/mujoco/loaders.py b/src/jaxsim/mujoco/loaders.py index 2d26c9b13..4a3012272 100644 --- a/src/jaxsim/mujoco/loaders.py +++ b/src/jaxsim/mujoco/loaders.py @@ -13,7 +13,17 @@ def load_rod_model( is_urdf: bool | None = None, model_name: str | None = None, ) -> rod.Model: - """""" + """ + Loads a ROD model from a URDF/SDF file or a ROD model. + + Args: + model_description: The URDF/SDF file or ROD model to load. + is_urdf: Whether the model description is a URDF file. + model_name: The name of the model to load from the resource. + + Returns: + rod.Model: The loaded ROD model. + """ # Parse the SDF resource. sdf_element = rod.Sdf.load(sdf=model_description, is_urdf=is_urdf) @@ -50,7 +60,15 @@ class RodModelToMjcf: def assets_from_rod_model( rod_model: rod.Model, ) -> dict[str, bytes]: - """""" + """ + Generates a dictionary of assets from a ROD model. + + Args: + rod_model: The ROD model to extract the assets from. + + Returns: + dict: A dictionary of assets. + """ import resolve_robotics_uri_py @@ -85,7 +103,17 @@ def add_floating_joint( base_link_name: str, floating_joint_name: str = "world_to_base", ) -> str: - """""" + """ + Adds a floating joint to a URDF string. + + Args: + urdf_string: The URDF string to modify. + base_link_name: The name of the base link to attach the floating joint. + floating_joint_name: The name of the floating joint to add. + + Returns: + str: The modified URDF string. + """ with tempfile.NamedTemporaryFile(mode="w+", suffix=".urdf") as urdf_file: @@ -132,7 +160,16 @@ def convert( plane_normal: tuple[float, float, float] = (0, 0, 1), heightmap: bool | None = None, ) -> tuple[str, dict[str, Any]]: - """""" + """ + Converts a ROD model to a Mujoco MJCF string. + + Args: + rod_model: The ROD model to convert. + considered_joints: The list of joint names to consider in the conversion. + + Returns: + tuple: A tuple containing the MJCF string and the assets dictionary. + """ # ------------------------------------- # Convert the model description to URDF @@ -468,7 +505,17 @@ def convert( plane_normal: tuple[float, float, float] = (0, 0, 1), heightmap: bool | None = None, ) -> tuple[str, dict[str, Any]]: - """""" + """ + Converts a URDF file to a Mujoco MJCF string. + + Args: + urdf: The URDF file to convert. + considered_joints: The list of joint names to consider in the conversion. + model_name: The name of the model to convert. + + Returns: + tuple: A tuple containing the MJCF string and the assets dictionary. + """ # Get the ROD model. rod_model = load_rod_model( @@ -495,7 +542,17 @@ def convert( plane_normal: tuple[float, float, float] = (0, 0, 1), heightmap: bool | None = None, ) -> tuple[str, dict[str, Any]]: - """""" + """ + Converts a SDF file to a Mujoco MJCF string. + + Args: + sdf: The SDF file to convert. + considered_joints: The list of joint names to consider in the conversion. + model_name: The name of the model to convert. + + Returns: + tuple: A tuple containing the MJCF string and the assets dictionary. + """ # Get the ROD model. rod_model = load_rod_model( diff --git a/src/jaxsim/mujoco/model.py b/src/jaxsim/mujoco/model.py index 845f6bb90..62d2912f8 100644 --- a/src/jaxsim/mujoco/model.py +++ b/src/jaxsim/mujoco/model.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import functools import pathlib from typing import Any, Callable @@ -18,7 +20,13 @@ class MujocoModelHelper: """ def __init__(self, model: mj.MjModel, data: mj.MjData | None = None) -> None: - """""" + """ + Initialize the MujocoModelHelper object. + + Args: + model: A Mujoco model object. + data: A Mujoco data object. If None, a new one will be created. + """ self.model = model self.data = data if data is not None else mj.MjData(self.model) @@ -34,8 +42,19 @@ def build_from_xml( mjcf_description: str | pathlib.Path, assets: dict[str, Any] = None, heightmap: HeightmapCallable | None = None, - ) -> "MujocoModelHelper": - """""" + ) -> MujocoModelHelper: + """ + Build a Mujoco model from an XML description and an optional assets dictionary. + + Args: + mjcf_description: A string containing the XML description of the Mujoco model + or a path to a file containing the XML description. + assets: An optional dictionary containing the assets of the model. + heightmap: A function in two variables that returns the height of a terrain + in the specified coordinate point. + Returns: + A MujocoModelHelper object. + """ # Read the XML description if it's a path to file mjcf_description = ( @@ -175,17 +194,17 @@ def is_dcm(R): # ================== def number_of_joints(self) -> int: - """""" + """Returns the number of joints in the model.""" return self.model.njnt def number_of_dofs(self) -> int: - """""" + """Returns the number of DoFs in the model.""" return self.model.nq def joint_names(self) -> list[str]: - """""" + """Returns the names of the joints in the model.""" return [ mj.mj_id2name(self.model, mj.mjtObj.mjOBJ_JOINT, idx) @@ -193,7 +212,7 @@ def joint_names(self) -> list[str]: ] def joint_dofs(self, joint_name: str) -> int: - """""" + """Returns the number of DoFs of a joint.""" if joint_name not in self.joint_names(): raise ValueError(f"Joint '{joint_name}' not found") @@ -201,7 +220,7 @@ def joint_dofs(self, joint_name: str) -> int: return self.data.joint(joint_name).qpos.size def joint_position(self, joint_name: str) -> npt.NDArray: - """""" + """Returns the position of a joint.""" if joint_name not in self.joint_names(): raise ValueError(f"Joint '{joint_name}' not found") @@ -209,7 +228,7 @@ def joint_position(self, joint_name: str) -> npt.NDArray: return self.data.joint(joint_name).qpos def joint_positions(self, joint_names: list[str] | None = None) -> npt.NDArray: - """""" + """Returns the positions of the joints.""" joint_names = joint_names if joint_names is not None else self.joint_names() @@ -220,7 +239,7 @@ def joint_positions(self, joint_names: list[str] | None = None) -> npt.NDArray: def set_joint_position( self, joint_name: str, position: npt.NDArray | float ) -> None: - """""" + """Sets the position of a joint.""" position = np.atleast_1d(np.array(position).squeeze()) @@ -239,7 +258,7 @@ def set_joint_position( def set_joint_positions( self, joint_names: list[str], positions: npt.NDArray | list[npt.NDArray] ) -> None: - """""" + """Set the positions of multiple joints.""" mask = self.mask_qpos(joint_names=tuple(joint_names)) self.data.qpos[mask] = positions @@ -249,12 +268,12 @@ def set_joint_positions( # ================== def number_of_bodies(self) -> int: - """""" + """Returns the number of bodies in the model.""" return self.model.nbody def body_names(self) -> list[str]: - """""" + """Returns the names of the bodies in the model.""" return [ mj.mj_id2name(self.model, mj.mjtObj.mjOBJ_BODY, idx) @@ -262,7 +281,7 @@ def body_names(self) -> list[str]: ] def body_position(self, body_name: str) -> npt.NDArray: - """""" + """Returns the position of a body.""" if body_name not in self.body_names(): raise ValueError(f"Body '{body_name}' not found") @@ -270,7 +289,7 @@ def body_position(self, body_name: str) -> npt.NDArray: return self.data.body(body_name).xpos def body_orientation(self, body_name: str, dcm: bool = False) -> npt.NDArray: - """""" + """Returns the orientation of a body.""" if body_name not in self.body_names(): raise ValueError(f"Body '{body_name}' not found") @@ -284,12 +303,12 @@ def body_orientation(self, body_name: str, dcm: bool = False) -> npt.NDArray: # ====================== def number_of_geometries(self) -> int: - """""" + """Returns the number of geometries in the model.""" return self.model.ngeom def geometry_names(self) -> list[str]: - """""" + """Returns the names of the geometries in the model.""" return [ mj.mj_id2name(self.model, mj.mjtObj.mjOBJ_GEOM, idx) @@ -297,7 +316,7 @@ def geometry_names(self) -> list[str]: ] def geometry_position(self, geometry_name: str) -> npt.NDArray: - """""" + """Returns the position of a geometry.""" if geometry_name not in self.geometry_names(): raise ValueError(f"Geometry '{geometry_name}' not found") @@ -307,7 +326,7 @@ def geometry_position(self, geometry_name: str) -> npt.NDArray: def geometry_orientation( self, geometry_name: str, dcm: bool = False ) -> npt.NDArray: - """""" + """Returns the orientation of a geometry.""" if geometry_name not in self.geometry_names(): raise ValueError(f"Geometry '{geometry_name}' not found") diff --git a/src/jaxsim/mujoco/visualizer.py b/src/jaxsim/mujoco/visualizer.py index 5cc587543..9af33631a 100644 --- a/src/jaxsim/mujoco/visualizer.py +++ b/src/jaxsim/mujoco/visualizer.py @@ -20,7 +20,17 @@ def __init__( height: int | None = None, **kwargs, ) -> None: - """""" + """ + Initialize the Mujoco video recorder. + + Args: + model: The Mujoco model. + data: The Mujoco data. + fps: The frames per second. + width: The width of the video. + height: The height of the video. + **kwargs: Additional arguments for the renderer. + """ width = width if width is not None else model.vis.global_.offwidth height = height if height is not None else model.vis.global_.offheight @@ -45,7 +55,7 @@ def __init__( def reset( self, model: mj.MjModel | None = None, data: mj.MjData | None = None ) -> None: - """""" + """Reset the model and data.""" self.frames = [] @@ -53,7 +63,7 @@ def reset( self.model = model if model is not None else self.model def render_frame(self, camera_name: str | None = None) -> npt.NDArray: - """""" + """Renders a frame.""" camera_name = camera_name or "track" mujoco.mj_forward(self.model, self.data) @@ -62,14 +72,14 @@ def render_frame(self, camera_name: str | None = None) -> npt.NDArray: return self.renderer.render() def record_frame(self, camera_name: str | None = None) -> None: - """""" + """Stores a frame in the buffer.""" camera_name = camera_name or "track" frame = self.render_frame(camera_name=camera_name) self.frames.append(frame) def write_video(self, path: pathlib.Path, exist_ok: bool = False) -> None: - """""" + """Writes the video to a file.""" if path.is_dir(): raise IsADirectoryError(f"The path '{path}' is a directory.") @@ -110,7 +120,13 @@ class MujocoVisualizer: def __init__( self, model: mj.MjModel | None = None, data: mj.MjData | None = None ) -> None: - """""" + """ + Initialize the Mujoco visualizer. + + Args: + model: The Mujoco model. + data: The Mujoco data. + """ self.data = data self.model = model @@ -121,7 +137,7 @@ def sync( model: mj.MjModel | None = None, data: mj.MjData | None = None, ) -> None: - """""" + """Updates the viewer with the current model and data.""" data = data if data is not None else self.data model = model if model is not None else self.model @@ -132,7 +148,7 @@ def sync( def open_viewer( self, model: mj.MjModel | None = None, data: mj.MjData | None = None ) -> mj.viewer.Handle: - """""" + """Opens a viewer.""" data = data if data is not None else self.data model = model if model is not None else self.model @@ -150,7 +166,7 @@ def open( data: mj.MjData | None = None, close_on_exit: bool = True, ) -> ContextManager[mujoco.viewer.Handle]: - """""" + """Context manager to open a viewer.""" handle = self.open_viewer(model=model, data=data) From 591f8b7fa90ca58dc5ff46129fccfad3ae9dd36a Mon Sep 17 00:00:00 2001 From: Filippo Luca Ferretti Date: Thu, 28 Mar 2024 11:45:03 +0100 Subject: [PATCH 12/21] Update index --- docs/index.rst | 91 +++++++++++++++++++++++++++++++------------------- 1 file changed, 56 insertions(+), 35 deletions(-) diff --git a/docs/index.rst b/docs/index.rst index 52c946397..7c16d0825 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -3,30 +3,68 @@ JAXsim A scalable physics engine and multibody dynamics library implemented with JAX. With JIT batteries 🔋 -.. warning:: - This project is still experimental, APIs could change without notice. - .. note:: This simulator currently focuses on locomotion applications. Only contacts with ground are supported. Features -------- -- Physics engine in reduced coordinates implemented with `JAX `_ in Python. -- JIT compilation of Python code for increased performance. -- Transparent support to execute logic on CPUs, GPUs, and TPUs. -- Parallel multi-body simulations on hardware accelerators for significantly increased throughput. -- Support for SDF models (and, upon conversion, URDF models). -- Collision detection between bodies and uneven ground surface. -- Soft contacts model supporting full friction cone and sticking / slipping transition. -- Complete support for inertial properties of rigid bodies. -- Revolute, prismatic, and fixed joints support. -- Integrators: forward Euler, semi-implicit Euler, Runge-Kutta 4. -- High-level classes for object-oriented programming. -- High-level classes to compute multi-body dynamics quantities from simulation state. -- High-level classes wrapping the low-level functional RBDAs with support of `multiple velocities representations `_. -- Default validation of JAX pytrees to prevent JIT re-compilations. -- Preliminary support for automatic differentiation of RBDAs. +.. grid:: + + .. grid-item:: + :columns: 12 12 12 6 + + .. card:: Performance + :class-card: sd-border-0 + :shadow: none + :class-title: sd-fs-5 + + .. div:: sd-font-normal + + Physics engine in reduced coordinates implemented with `JAX `_. + Compatibility with JIT compilation for increased performance and transparent support to execute logic on CPUs, GPUs, and TPUs. + Parallel multi-body simulations on hardware accelerators for significantly increased throughput + + .. grid-item:: + :columns: 12 12 12 6 + + .. card:: Model Parsing + :class-card: sd-border-0 + :shadow: none + :class-title: sd-fs-5 + + .. div:: sd-font-normal + + Support for SDF models (and, upon conversion, URDF models). Revolute, prismatic, and fixed joints supported. + + .. grid-item:: + :columns: 12 12 12 6 + + .. card:: Automatic Differentiation + :class-card: sd-border-0 + :shadow: none + :class-title: sd-fs-5 + + .. div:: sd-font-normal + + Support for automatic differentiation of rigid body dynamics algorithms (RBDAs) for model-based robotics research. + Soft contacts model supporting full friction cone and sticking / slipping transition. + + .. grid-item:: + :columns: 12 12 12 6 + + .. card:: Complex Dynamics + :class-card: sd-border-0 + :shadow: none + :class-title: sd-fs-5 + + .. div:: sd-font-normal + + JAXsim provides a variety of integrators for the simulation of multibody dynamics, including RK4, Heun, Euler, and more. + Support of `multiple velocities representations `_. + + +---- .. toctree:: :hidden: @@ -73,23 +111,6 @@ Here below we summarize the differences between the projects: - Contrarily to brax, JAXsim only supports collision detection between bodies and a compliant ground surface. - The RBDAs of JAXsim support automatic differentiation, but this functionality has not been thoroughly tested. -Contributing ------------- - -Pull requests are welcome. -For major changes, please open an issue first to discuss what you would like to change. - -Citing ------- - -.. code-block:: bibtex - - @software{ferigo_jaxsim_2022, - author = {Diego Ferigo and Silvio Traversaro and Daniele Pucci}, - title = {{JAXsim}: A Physics Engine in Reduced Coordinates and Multibody Dynamics Library for Control and Robot Learning}, - url = {http://github.com/ami-iit/jaxsim}, - year = {2022}, - } People ------ From 6e97c43a38bdede8d0d394221a2389ebac00fba6 Mon Sep 17 00:00:00 2001 From: Filippo Luca Ferretti Date: Thu, 28 Mar 2024 13:30:24 +0100 Subject: [PATCH 13/21] Add missing sphinx-design dependency --- docs/conf.py | 1 + environment.yml | 1 + 2 files changed, 2 insertions(+) diff --git a/docs/conf.py b/docs/conf.py index 7a4eb3c30..e1ce4e33d 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -58,6 +58,7 @@ def _recursive_add_annotations_import(): "sphinx_autodoc_typehints", "sphinx_multiversion", "enum_tools.autoenum", + "sphinx_design", ] # -- Options for intersphinx extension diff --git a/environment.yml b/environment.yml index e182c4c62..17e2f272a 100644 --- a/environment.yml +++ b/environment.yml @@ -38,6 +38,7 @@ dependencies: - sphinx - sphinx-autodoc-typehints - sphinx-copybutton + - sphinx-design - sphinx_fontawesome - sphinx-jinja2-compat - sphinx-multiversion From 6431b8f326e2dc6a3e567cecf3b05125ad7c7ebe Mon Sep 17 00:00:00 2001 From: Filippo Luca Ferretti Date: Thu, 28 Mar 2024 13:30:36 +0100 Subject: [PATCH 14/21] Add missing module --- docs/index.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/index.rst b/docs/index.rst index 7c16d0825..236bafea6 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -85,7 +85,7 @@ Features modules/mujoco modules/parsers modules/utils - + modules/typing Examples -------- From e0bcc44bcd747950281989674df82ddd2f27f2cb Mon Sep 17 00:00:00 2001 From: Filippo Luca Ferretti Date: Thu, 28 Mar 2024 13:31:46 +0100 Subject: [PATCH 15/21] Remove deprecated math modules --- docs/modules/math.rst | 16 ++-------------- 1 file changed, 2 insertions(+), 14 deletions(-) diff --git a/docs/modules/math.rst b/docs/modules/math.rst index 0a7e8de3f..bc774f985 100644 --- a/docs/modules/math.rst +++ b/docs/modules/math.rst @@ -1,13 +1,9 @@ -Math +Math ==== .. currentmodule:: jaxsim.math - -.. automodule:: jaxsim.math.adjoint - :members: - :undoc-members: -.. automodule:: jaxsim.math.conv +.. automodule:: jaxsim.math.adjoint :members: :undoc-members: @@ -19,14 +15,6 @@ Math :members: :undoc-members: -.. automodule:: jaxsim.math.joint - :members: - :undoc-members: - -.. automodule:: jaxsim.math.plucker - :members: - :undoc-members: - .. automodule:: jaxsim.math.quaternion :members: :undoc-members: From 481b1facd9184003e6c1dc6c9aa104a51edd641c Mon Sep 17 00:00:00 2001 From: Filippo Luca Ferretti Date: Thu, 28 Mar 2024 13:32:14 +0100 Subject: [PATCH 16/21] Add missing raw-string identifiers --- src/jaxsim/api/model.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/jaxsim/api/model.py b/src/jaxsim/api/model.py index 1d46272a2..94119ded6 100644 --- a/src/jaxsim/api/model.py +++ b/src/jaxsim/api/model.py @@ -986,8 +986,8 @@ def free_floating_gravity_forces( def free_floating_bias_forces( model: JaxSimModel, data: js.data.JaxSimModelData ) -> jtp.Vector: - """ - Compute the free-floating bias forces :math:`h(\\mathbf{q}, \boldsymbol{\nu})` + r""" + Compute the free-floating bias forces :math:`h(\mathbf{q}, \boldsymbol{\nu})` of the model. Args: From abd6b12b8366de38bb3306467a2eb91d5157b520 Mon Sep 17 00:00:00 2001 From: Filippo Luca Ferretti Date: Thu, 28 Mar 2024 15:06:22 +0100 Subject: [PATCH 17/21] Add autosummary for complex modules --- docs/index.rst | 11 ++--------- docs/modules/api.rst | 5 ++--- docs/modules/index.rst | 19 +++++++++++++++++++ docs/modules/integrators.rst | 1 + docs/modules/rbda.rst | 10 ++++++++++ 5 files changed, 34 insertions(+), 12 deletions(-) create mode 100644 docs/modules/index.rst diff --git a/docs/index.rst b/docs/index.rst index 236bafea6..2a0cefd97 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -75,17 +75,10 @@ Features .. toctree:: :hidden: - :maxdepth: 2 + :maxdepth: 1 :caption: JAXsim API - modules/api - modules/integrators - modules/rbda - modules/math - modules/mujoco - modules/parsers - modules/utils - modules/typing + modules/index Examples -------- diff --git a/docs/modules/api.rst b/docs/modules/api.rst index 3f79d76f4..e162063d0 100644 --- a/docs/modules/api.rst +++ b/docs/modules/api.rst @@ -1,9 +1,8 @@ -API -========== +Functional API +============== .. currentmodule:: jaxsim.api - Model ~~~~~ diff --git a/docs/modules/index.rst b/docs/modules/index.rst new file mode 100644 index 000000000..b8ab64b16 --- /dev/null +++ b/docs/modules/index.rst @@ -0,0 +1,19 @@ +Public API +========== + +.. _modules: + +Modules +------- + +.. toctree:: + :maxdepth: 1 + + api + integrators + math + mujoco + parsers + rbda + typing + utils diff --git a/docs/modules/integrators.rst b/docs/modules/integrators.rst index 2a1e5327d..c9e776e81 100644 --- a/docs/modules/integrators.rst +++ b/docs/modules/integrators.rst @@ -3,6 +3,7 @@ Integrators .. currentmodule:: jaxsim.integrators + Common ~~~~~~ diff --git a/docs/modules/rbda.rst b/docs/modules/rbda.rst index 60e69be97..dac16010f 100644 --- a/docs/modules/rbda.rst +++ b/docs/modules/rbda.rst @@ -5,6 +5,16 @@ This module provides a set of algorithms for rigid body dynamics. .. currentmodule:: jaxsim.rbda +.. autosummary:: + :toctree: _autosummary + + aba + collidable_points + crba + forward_kinematics + jacobian + soft_contacts + utils Articulated Body Algorithm ~~~~~~~~~~~~~~~~~~~~~~~~~~ From fdff2018df54568f937526bc5b1c6d557295ea4d Mon Sep 17 00:00:00 2001 From: Filippo Luca Ferretti Date: Thu, 28 Mar 2024 16:07:31 +0100 Subject: [PATCH 18/21] Disable forward referenced annotations --- docs/conf.py | 18 +----------------- 1 file changed, 1 insertion(+), 17 deletions(-) diff --git a/docs/conf.py b/docs/conf.py index e1ce4e33d..29f674cbe 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -4,23 +4,7 @@ from pkg_resources import get_distribution - -def _add_annotations_import(path): - with open(path, "r+") as f: - contents = f.read() - if not contents.startswith("from __future__ import annotations"): - f.seek(0, 0) - f.write("from __future__ import annotations " + contents) - - -def _recursive_add_annotations_import(): - for path, _, files in os.walk("../jaxsim/"): - for file in [f for f in files if f.endswith(".py")]: - _add_annotations_import(os.path.join(path, file)) - - -if "READTHEDOCS" in os.environ: - _recursive_add_annotations_import() +import jaxsim # -- Version information From cdc8f3d42bb989a062de809c55957d59cdbf265f Mon Sep 17 00:00:00 2001 From: Filippo Luca Ferretti Date: Tue, 2 Apr 2024 12:59:13 +0200 Subject: [PATCH 19/21] Lint --- .github/workflows/ci_cd.yml | 2 +- LICENSE | 2 +- docs/guide/install.rst | 4 +- docs/modules/parsers.rst | 2 +- docs/modules/typing.rst | 4 +- examples/.gitattributes | 1 - examples/.gitignore | 1 - examples/README.md | 6 +- examples/assets/cartpole.urdf | 156 +++++++++++++++++----------------- pyproject.toml | 1 + src/jaxsim/api/model.py | 2 +- 11 files changed, 90 insertions(+), 91 deletions(-) diff --git a/.github/workflows/ci_cd.yml b/.github/workflows/ci_cd.yml index eeb3c1154..bbb24d0bb 100644 --- a/.github/workflows/ci_cd.yml +++ b/.github/workflows/ci_cd.yml @@ -92,7 +92,7 @@ jobs: - name: Document installed pip packages shell: bash run: pip list --verbose - + - name: Import the package run: python -c "import jaxsim" diff --git a/LICENSE b/LICENSE index 3c486a482..2755199ec 100644 --- a/LICENSE +++ b/LICENSE @@ -1,6 +1,6 @@ BSD 3-Clause License -Copyright (c) 2022, Artificial and Mechanical Intelligence +Copyright (c) 2022, Artificial and Mechanical Intelligence All rights reserved. Redistribution and use in source and binary forms, with or without diff --git a/docs/guide/install.rst b/docs/guide/install.rst index 695343f46..a83c0d981 100644 --- a/docs/guide/install.rst +++ b/docs/guide/install.rst @@ -6,12 +6,12 @@ Installation Prerequisites ------------- -JAXsim requires Python 3.11 or later. +JAXsim requires Python 3.11 or later. Basic Installation ------------------ -You can install the project with using `conda`_: +You can install the project with using `conda`_: .. code-block:: bash diff --git a/docs/modules/parsers.rst b/docs/modules/parsers.rst index 367d21d19..f79ce1183 100644 --- a/docs/modules/parsers.rst +++ b/docs/modules/parsers.rst @@ -1,4 +1,4 @@ -Parsers +Parsers ======= .. automodule:: jaxsim.parsers.descriptions.collision diff --git a/docs/modules/typing.rst b/docs/modules/typing.rst index 65a3fe78f..d7bb370bf 100644 --- a/docs/modules/typing.rst +++ b/docs/modules/typing.rst @@ -1,4 +1,4 @@ -Typing +Typing ====== .. currentmodule:: jaxsim.typing @@ -14,4 +14,4 @@ Typing IntJax ArrayJax VectorJax - MatrixJax \ No newline at end of file + MatrixJax diff --git a/examples/.gitattributes b/examples/.gitattributes index 16ef5c5f7..d5799bd69 100644 --- a/examples/.gitattributes +++ b/examples/.gitattributes @@ -1,3 +1,2 @@ # GitHub syntax highlighting pixi.lock linguist-language=YAML - diff --git a/examples/.gitignore b/examples/.gitignore index 84ef293bb..c9314b7c2 100644 --- a/examples/.gitignore +++ b/examples/.gitignore @@ -1,3 +1,2 @@ # pixi environments .pixi - diff --git a/examples/README.md b/examples/README.md index 2fee0bbdb..ced417a98 100644 --- a/examples/README.md +++ b/examples/README.md @@ -21,12 +21,12 @@ The simplest way to run the examples is by accessing the provided Google Colab n For local execution, follow these steps: -1. **Install `pixi`:** +1. **Install `pixi`:** As per the [official documentation](https://pixi.sh/#installation): ```bash -curl -fsSL https://pixi.sh/install.sh | bash +curl -fsSL https://pixi.sh/install.sh | bash ``` 2. **Run the Example Notebook:** @@ -37,4 +37,4 @@ Use `pixi run ` to execute the example notebook locally, e.g.: pixi run PD_controller ``` -This command will automatically handle the installation of necessary dependencies and execute the examples within a self-contained environment \ No newline at end of file +This command will automatically handle the installation of necessary dependencies and execute the examples within a self-contained environment diff --git a/examples/assets/cartpole.urdf b/examples/assets/cartpole.urdf index f58e2987e..8b8c9dc93 100644 --- a/examples/assets/cartpole.urdf +++ b/examples/assets/cartpole.urdf @@ -1,83 +1,83 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + - + - \ No newline at end of file + diff --git a/pyproject.toml b/pyproject.toml index d3b070f38..b357f63b3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -66,3 +66,4 @@ select = [ "**/{tests,docs,tools}/*" = ["E402"] "**/{tests}/*" = ["B007"] "__init__.py" = ["F401"] +"docs/conf.py" = ["F401"] diff --git a/src/jaxsim/api/model.py b/src/jaxsim/api/model.py index 94119ded6..e250678f3 100644 --- a/src/jaxsim/api/model.py +++ b/src/jaxsim/api/model.py @@ -986,7 +986,7 @@ def free_floating_gravity_forces( def free_floating_bias_forces( model: JaxSimModel, data: js.data.JaxSimModelData ) -> jtp.Vector: - r""" + r""" Compute the free-floating bias forces :math:`h(\mathbf{q}, \boldsymbol{\nu})` of the model. From b5c30c99b9a76ea290e47b4bf94f0572d1bd8531 Mon Sep 17 00:00:00 2001 From: Filippo Luca Ferretti Date: Wed, 3 Apr 2024 10:47:21 +0200 Subject: [PATCH 20/21] Add missing dependency on docs environment --- environment.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/environment.yml b/environment.yml index 17e2f272a..b172f8a85 100644 --- a/environment.yml +++ b/environment.yml @@ -45,3 +45,4 @@ dependencies: - sphinx_rtd_theme - sphinx-book-theme - sphinx-toolbox + - pip From 34b63b2295b60095ffa79378ca13a033a68c3446 Mon Sep 17 00:00:00 2001 From: Filippo Luca Ferretti Date: Wed, 3 Apr 2024 11:04:48 +0200 Subject: [PATCH 21/21] Set missing environment variable --- docs/conf.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/docs/conf.py b/docs/conf.py index 29f674cbe..fc03564dc 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -4,6 +4,12 @@ from pkg_resources import get_distribution +if os.environ.get("READTHEDOCS"): + checkout_name = os.path.basename(os.path.dirname(os.path.realpath(__file__))) + os.environ["CONDA_PREFIX"] = os.path.realpath( + os.path.join("..", "..", "conda", checkout_name) + ) + import jaxsim # -- Version information