jaxsim 0.4.3.dev245__tar.gz → 0.4.3.dev271__tar.gz

{jaxsim-0.4.3.dev245 → jaxsim-0.4.3.dev271}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.3.dev245
+Version: 0.4.3.dev271
 Summary: A differentiable physics engine and multibody dynamics library for control and robot learning.
 Author-email: Diego Ferigo <dgferigo@gmail.com>
 Maintainer-email: Diego Ferigo <dgferigo@gmail.com>, Filippo Luca Ferretti <filippo.ferretti@iit.it>

{jaxsim-0.4.3.dev245 → jaxsim-0.4.3.dev271}/examples/jaxsim_as_physics_engine.ipynb

@@ -63,8 +63,9 @@
     "import jax\n",
     "import jax.numpy as jnp\n",
     "import jaxsim.api as js\n",
+    "import jaxsim\n",
     "import rod\n",
-    "from jaxsim import integrators, logging\n",
+    "from jaxsim import logging\n",
     "from rod.builder.primitives import SphereBuilder\n",
     "\n",
     "logging.set_logging_level(logging.LoggingLevel.WARNING)\n",
@@ -142,8 +143,11 @@
     "\n",
     "- `model`: an object that defines the dynamics of the system.\n",
     "- `data`: an object that contains the state of the system.\n",
-    "- `integrator`: an object that defines the integration method.\n",
-    "- `integrator_state`: an object that contains the state of the integrator."
+    "- `integrator` *(Optional)*: an object that defines the integration method.\n",
+    "- `integrator_state` *(Optional)*: an object that contains the state of the integrator.\n",
+    "\n",
+    "The `JaxSimModel` object contains the simulation time step, the integrator and the contact model.\n",
+    "In this example, we will explicity pass an integrator class to the `model` object and we will use the default `SoftContacts` contact model."
    ]
   },
   {
@@ -157,7 +161,9 @@
     "# Create the JaxSim model.\n",
     "# This is shared among all the parallel instances.\n",
     "model = js.model.JaxSimModel.build_from_model_description(\n",
-    "    model_description=model_sdf_string, time_step=0.001\n",
+    "    model_description=model_sdf_string,\n",
+    "    time_step=0.001,\n",
+    "    integrator=jaxsim.integrators.fixed_step.Heun2,\n",
     ")\n",
     "\n",
     "# Create the data of a single model.\n",
@@ -240,21 +246,7 @@
    },
    "outputs": [],
    "source": [
-    "# Create the integrator.\n",
-    "integrator = integrators.fixed_step.Heun2SO3.build(\n",
-    "    dynamics=js.ode.wrap_system_dynamics_for_integration(\n",
-    "        model=model,\n",
-    "        data=data_single,\n",
-    "        system_dynamics=js.ode.system_dynamics,\n",
-    "    ),\n",
-    ")\n",
-    "\n",
-    "# Initialize the integrator.\n",
-    "integrator_state = integrator.init(\n",
-    "    x0=data_single.state,\n",
-    "    t0=0.0,\n",
-    "    dt=model.time_step,\n",
-    ")\n",
+    "print(f\"Using integrator: {model.integrator}\")\n",
     "\n",
     "# Initialize the simulated time.\n",
     "T = jnp.arange(start=0, stop=1.0, step=model.time_step)"
@@ -324,46 +316,42 @@
     "def step_single(\n",
     "    model: js.model.JaxSimModel,\n",
     "    data: js.data.JaxSimModelData,\n",
-    "    integrator_state: dict[str, Any],\n",
     ") -> tuple[js.data.JaxSimModelData, dict[str, Any]]:\n",
     "\n",
     "    # Close step over static arguments.\n",
     "    return js.model.step(\n",
     "        model=model,\n",
     "        data=data,\n",
-    "        integrator=integrator,\n",
-    "        integrator_state=integrator_state,\n",
     "        link_forces=None,\n",
     "        joint_force_references=None,\n",
     "    )\n",
     "\n",
     "\n",
     "@jax.jit\n",
-    "@functools.partial(jax.vmap, in_axes=(None, 0, None))\n",
+    "@functools.partial(jax.vmap, in_axes=(None, 0))\n",
     "def step_parallel(\n",
     "    model: js.model.JaxSimModel,\n",
     "    data: js.data.JaxSimModelData,\n",
-    "    integrator_state: dict[str, Any],\n",
     ") -> tuple[js.data.JaxSimModelData, dict[str, Any]]:\n",
     "\n",
     "    return step_single(\n",
-    "        model=model, data=data, integrator_state=integrator_state\n",
+    "        model=model, data=data\n",
     "    )\n",
     "\n",
     "\n",
     "# The first run will be slow since JAX needs to JIT-compile the functions.\n",
-    "_ = step_single(model, data_single, integrator_state)\n",
-    "_ = step_parallel(model, data_batch_t0, integrator_state)\n",
+    "_ = step_single(model, data_single)\n",
+    "_ = step_parallel(model, data_batch_t0)\n",
     "\n",
     "# Benchmark the execution of a single step.\n",
     "print(\"\\nSingle simulation step:\")\n",
-    "%timeit step_single(model, data_single, integrator_state)\n",
+    "%timeit step_single(model, data_single)\n",
     "\n",
     "# On hardware accelerators, there's a range of batch_size values where\n",
     "# increasing the number of parallel instances doesn't affect computation time.\n",
     "# This range depends on the GPU/TPU specifications.\n",
     "print(f\"\\nParallel simulation steps (batch_size={batch_size} on {jax.devices()[0]}):\")\n",
-    "%timeit step_parallel(model, data_batch_t0, integrator_state)"
+    "%timeit step_parallel(model, data_batch_t0)"
    ]
   },
   {
@@ -381,7 +369,7 @@
     "\n",
     "for _ in T:\n",
     "\n",
-    "    data, integrator_state = step_parallel(model, data, integrator_state)\n",
+    "    data, _ = step_parallel(model, data)\n",
     "    data_trajectory_list.append(data)"
    ]
   },
@@ -404,9 +392,7 @@
    "source": [
     "# Convert a list of PyTrees to a batched PyTree.\n",
     "# This operation is called 'tree transpose' in JAX.\n",
-    "data_trajectory = jax.tree.map(\n",
-    "    lambda *leafs: jnp.stack(leafs), *data_trajectory_list\n",
-    ")\n",
+    "data_trajectory = jax.tree.map(lambda *leafs: jnp.stack(leafs), *data_trajectory_list)\n",
     "\n",
     "print(f\"W_p_B: shape={data_trajectory.base_position().shape}\")"
    ]
@@ -448,6 +434,11 @@
     "\n",
     "Have fun!"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": []
   }
  ],
  "metadata": {
@@ -459,7 +450,8 @@
    "toc_visible": true
   },
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "jaxsim",
+   "language": "python",
    "name": "python3"
   },
   "language_info": {
@@ -472,7 +464,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.8"
+   "version": "3.12.7"
   }
  },
  "nbformat": 4,

{jaxsim-0.4.3.dev245 → jaxsim-0.4.3.dev271}/examples/jaxsim_for_robot_controllers.ipynb

@@ -123,7 +123,6 @@
     "# @title Create the model and its data\n",
     "\n",
     "import jaxsim.api as js\n",
-    "from jaxsim import integrators\n",
     "\n",
     "# Create the model from the model description.\n",
     "model = js.model.JaxSimModel.build_from_model_description(\n",
@@ -143,23 +142,7 @@
    },
    "outputs": [],
    "source": [
-    "# @title Select the integrator\n",
-    "\n",
-    "# Create the integrator.\n",
-    "integrator = integrators.fixed_step.RungeKutta4.build(\n",
-    "    dynamics=js.ode.wrap_system_dynamics_for_integration(\n",
-    "        model=model,\n",
-    "        data=data_zero,\n",
-    "        system_dynamics=js.ode.system_dynamics,\n",
-    "    ),\n",
-    ")\n",
-    "\n",
-    "# Initialize the integrator.\n",
-    "integrator_state = integrator.init(\n",
-    "    x0=data_zero.state,\n",
-    "    t0=0.0,\n",
-    "    dt=model.time_step,\n",
-    ")\n",
+    "# @title Define simulation parameters\n",
     "\n",
     "# Initialize the simulated time.\n",
     "T = jnp.arange(start=0, stop=5.0, step=model.time_step)"
@@ -255,11 +238,9 @@
     "for _ in T:\n",
     "\n",
     "    # Step the JaxSim simulation.\n",
-    "    data, integrator_state = js.model.step(\n",
+    "    data, _ = js.model.step(\n",
     "        model=model,\n",
     "        data=data,\n",
-    "        integrator=integrator,\n",
-    "        integrator_state=integrator_state,\n",
     "        joint_force_references=None,\n",
     "        link_forces=None,\n",
     "    )\n",
@@ -359,7 +340,7 @@
     "    ṡ = data.joint_velocities()\n",
     "\n",
     "    # Compute the actuated joint torques.\n",
-    "    s_star = - kp * (s - s_des) - kd * (ṡ - s_dot_des)\n",
+    "    s_star = -kp * (s - s_des) - kd * (ṡ - s_dot_des)\n",
     "    τ = Mss @ s_star + hs\n",
     "\n",
     "    return τ"
@@ -407,11 +388,9 @@
     "    )\n",
     "\n",
     "    # Step the JaxSim simulation.\n",
-    "    data, integrator_state = js.model.step(\n",
+    "    data, _ = js.model.step(\n",
     "        model=model,\n",
     "        data=data,\n",
-    "        integrator=integrator,\n",
-    "        integrator_state=integrator_state,\n",
     "        joint_force_references=τ,\n",
     "    )\n",
     "\n",
@@ -461,7 +440,8 @@
    "toc_visible": true
   },
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "jaxsim",
+   "language": "python",
    "name": "python3"
   },
   "language_info": {
@@ -474,7 +454,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.8"
+   "version": "3.12.7"
   }
  },
  "nbformat": 4,

{jaxsim-0.4.3.dev245 → jaxsim-0.4.3.dev271}/src/jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.4.3.dev245'
-__version_tuple__ = version_tuple = (0, 4, 3, 'dev245')
+__version__ = version = '0.4.3.dev271'
+__version_tuple__ = version_tuple = (0, 4, 3, 'dev271')

{jaxsim-0.4.3.dev245 → jaxsim-0.4.3.dev271}/src/jaxsim/api/model.py

@@ -54,6 +54,10 @@ class JaxSimModel(JaxsimDataclass):
         default=None, repr=False
     )
 
+    integrator: Static[jaxsim.integrators.Integrator | None] = dataclasses.field(
+        default=None, repr=False
+    )
+
     _description: Static[wrappers.HashlessObject[ModelDescription | None]] = (
         dataclasses.field(default=None, repr=False)
     )
@@ -93,12 +97,16 @@ class JaxSimModel(JaxsimDataclass):
     # Initialization and state
     # ========================
 
-    @staticmethod
+    @classmethod
     def build_from_model_description(
+        cls,
         model_description: str | pathlib.Path | rod.Model,
-        model_name: str | None = None,
         *,
+        model_name: str | None = None,
         time_step: jtp.FloatLike | None = None,
+        integrator: (
+            jaxsim.integrators.Integrator | type[jaxsim.integrators.Integrator] | None
+        ) = None,
         terrain: jaxsim.terrain.Terrain | None = None,
         contact_model: jaxsim.rbda.contacts.ContactModel | None = None,
         is_urdf: bool | None = None,
@@ -120,6 +128,10 @@ class JaxSimModel(JaxsimDataclass):
             contact_model:
                 The contact model to consider.
                 If not specified, a soft contacts model is used.
+            integrator:
+                The integrator to use. If not specified, a default one is used.
+                This argument can either be a pre-built integrator instance or one
+                of the integrator classes defined in JaxSim.
             is_urdf:
                 The optional flag to force the model description to be parsed as a URDF.
                 This is usually automatically inferred.
@@ -146,10 +158,11 @@ class JaxSimModel(JaxsimDataclass):
             )
 
         # Build the model.
-        model = JaxSimModel.build(
+        model = cls.build(
             model_description=intermediate_description,
             model_name=model_name,
             time_step=time_step,
+            integrator=integrator,
             terrain=terrain,
             contact_model=contact_model,
         )
@@ -160,12 +173,16 @@ class JaxSimModel(JaxsimDataclass):
 
         return model
 
-    @staticmethod
+    @classmethod
     def build(
+        cls,
         model_description: ModelDescription,
-        model_name: str | None = None,
         *,
+        model_name: str | None = None,
         time_step: jtp.FloatLike | None = None,
+        integrator: (
+            jaxsim.integrators.Integrator | type[jaxsim.integrators.Integrator] | None
+        ) = None,
         terrain: jaxsim.terrain.Terrain | None = None,
         contact_model: jaxsim.rbda.contacts.ContactModel | None = None,
     ) -> JaxSimModel:
@@ -182,6 +199,11 @@ class JaxSimModel(JaxsimDataclass):
                 The default time step to consider for the simulation. It can be
                 manually overridden in the function that steps the simulation.
             terrain: The terrain to consider (the default is a flat infinite plane).
+                The optional name of the model overriding the physics model name.
+            integrator:
+                The integrator to use. If not specified, a default one is used.
+                This argument can either be a pre-built integrator instance or one
+                of the integrator classes defined in JaxSim.
             contact_model:
                 The contact model to consider.
                 If not specified, a soft contacts model is used.
@@ -195,23 +217,62 @@ class JaxSimModel(JaxsimDataclass):
 
         # Consider the default terrain (a flat infinite plane) if not specified.
         terrain = (
-            terrain or JaxSimModel.__dataclass_fields__["terrain"].default_factory()
+            terrain
+            if terrain is not None
+            else JaxSimModel.__dataclass_fields__["terrain"].default_factory()
         )
 
         # Consider the default time step if not specified.
         time_step = (
-            time_step or JaxSimModel.__dataclass_fields__["time_step"].default_factory()
+            time_step
+            if time_step is not None
+            else JaxSimModel.__dataclass_fields__["time_step"].default_factory()
         )
 
         # Create the default contact model.
         # It will be populated with an initial estimation of good parameters.
         # While these might not be the best, they are a good starting point.
-        contact_model = contact_model or jaxsim.rbda.contacts.SoftContacts.build(
-            terrain=terrain, parameters=None
+        contact_model = (
+            contact_model
+            if contact_model is not None
+            else jaxsim.rbda.contacts.SoftContacts.build(
+                terrain=terrain, parameters=None
+            )
         )
 
+        # Build the integrator if not provided.
+        match integrator:
+
+            # If None, build a default integrator.
+            case None:
+
+                integrator = jaxsim.integrators.fixed_step.Heun2SO3.build(
+                    dynamics=js.ode.wrap_system_dynamics_for_integration(
+                        system_dynamics=js.ode.system_dynamics
+                    )
+                )
+
+            # If it's a pre-built integrator (also a custom one from the user)
+            # just use it as is.
+            case _ if isinstance(integrator, jaxsim.integrators.Integrator):
+                pass
+
+            # If an integrator class is passed, assume that it is a JaxSim integrator
+            # and build it with the default system dynamics.
+            case _ if issubclass(integrator, jaxsim.integrators.Integrator):
+
+                integrator_cls = integrator
+                integrator = integrator_cls.build(
+                    dynamics=js.ode.wrap_system_dynamics_for_integration(
+                        system_dynamics=js.ode.system_dynamics
+                    )
+                )
+
+            case _:
+                raise ValueError(f"Invalid integrator: {integrator}")
+
         # Build the model.
-        model = JaxSimModel(
+        model = cls(
             model_name=model_name,
             kin_dyn_parameters=js.kin_dyn_parameters.KynDynParameters.build(
                 model_description=model_description
@@ -219,6 +280,7 @@ class JaxSimModel(JaxsimDataclass):
             time_step=time_step,
             terrain=terrain,
             contact_model=contact_model,
+            integrator=integrator,
             # The following is wrapped as hashless since it's a static argument, and we
             # don't want to trigger recompilation if it changes. All relevant parameters
             # needed to compute kinematics and dynamics quantities are stored in the
@@ -404,6 +466,7 @@ def reduce(
     reduced_model = JaxSimModel.build(
         model_description=reduced_intermediate_description,
         model_name=model.name(),
+        time_step=model.time_step,
         terrain=model.terrain,
         contact_model=model.contact_model,
     )
@@ -1912,10 +1975,10 @@ def step(
     model: JaxSimModel,
     data: js.data.JaxSimModelData,
     *,
-    integrator: jaxsim.integrators.Integrator,
     t0: jtp.FloatLike = 0.0,
     dt: jtp.FloatLike | None = None,
-    integrator_state: dict[str, Any] | None = None,
+    integrator: jaxsim.integrators.Integrator | None = None,
+    integrator_metadata: dict[str, Any] | None = None,
     link_forces: jtp.MatrixLike | None = None,
     joint_force_references: jtp.VectorLike | None = None,
     **kwargs,
@@ -1927,7 +1990,7 @@ def step(
         model: The model to consider.
         data: The data of the considered model.
         integrator: The integrator to use.
-        integrator_state: The state of the integrator.
+        integrator_metadata: The metadata of the integrator, if needed.
         t0: The initial time to consider. Only relevant for time-dependent dynamics.
         dt: The time step to consider. If not specified, it is read from the model.
         link_forces:
@@ -1937,8 +2000,9 @@ def step(
         kwargs: Additional kwargs to pass to the integrator.
 
     Returns:
-        A tuple containing the new data of the model
-        and the new state of the integrator.
+        A tuple containing the new data of the model and a dictionary of auxiliary
+        data computed during the step. If the integrator has metadata, the dictionary
+        will contain the new metadata stored in the `integrator_metadata` key.
 
     Note:
         In order to reduce the occurrences of frame conversions performed internally,
@@ -1953,8 +2017,9 @@ def step(
     integrator_kwargs = kwargs.pop("integrator_kwargs", {})
     integrator_kwargs = kwargs | integrator_kwargs
 
-    # Initialize the integrator state.
-    integrator_state_t0 = integrator_state if integrator_state is not None else dict()
+    # Extract the integrator and the optional metadata.
+    integrator_metadata_t0 = integrator_metadata
+    integrator = integrator if integrator is not None else model.integrator
 
     # Initialize the time-related variables.
     state_t0 = data.state
@@ -2010,11 +2075,11 @@ def step(
         τ_references = references.joint_force_references(model=model)
 
     # Step the dynamics forward.
-    state_tf, integrator_state_tf = integrator.step(
+    state_tf, integrator_metadata_tf = integrator.step(
         x0=state_t0,
         t0=t0,
         dt=dt,
-        params=integrator_state_t0,
+        metadata=integrator_metadata_t0,
         # Always inject the current (model, data) pair into the system dynamics
         # considered by the integrator, and include the input variables represented
         # by the pair (f_L, τ_references).
@@ -2091,13 +2156,17 @@ def step(
                     )
                 )
 
-            # Reset the generalized velocity.
-            data_tf = data_tf.reset_base_velocity(BW_nu_post_impact[0:6])
-            data_tf = data_tf.reset_joint_velocities(BW_nu_post_impact[6:])
+                # Reset the generalized velocity.
+                data_tf = data_tf.reset_base_velocity(BW_nu_post_impact[0:6])
+                data_tf = data_tf.reset_joint_velocities(BW_nu_post_impact[6:])
 
     # Restore the input velocity representation.
     data_tf = data_tf.replace(
         velocity_representation=data.velocity_representation, validate=False
     )
 
-    return data_tf, integrator_state_tf
+    return data_tf, {} | (
+        dict(integrator_metadata=integrator_metadata_tf)
+        if integrator_metadata is not None
+        else {}
+    )

{jaxsim-0.4.3.dev245 → jaxsim-0.4.3.dev271}/src/jaxsim/api/ode.py

@@ -24,41 +24,45 @@ class SystemDynamicsFromModelAndData(Protocol):
 
 
 def wrap_system_dynamics_for_integration(
-    model: js.model.JaxSimModel,
-    data: js.data.JaxSimModelData,
     *,
     system_dynamics: SystemDynamicsFromModelAndData,
-    **kwargs,
+    **kwargs: dict[str, Any],
 ) -> jaxsim.integrators.common.SystemDynamics[ODEState, ODEState]:
     """
-    Wrap generic system dynamics operating on `JaxSimModel` and `JaxSimModelData`
-    for integration with `jaxsim.integrators`.
+    Wrap the system dynamics considered by JaxSim integrators in a generic
+    `f(x, t, **u, **parameters)` function.
 
     Args:
-        model: The model to consider.
-        data: The data of the considered model.
         system_dynamics: The system dynamics to wrap.
         **kwargs: Additional kwargs to close over the system dynamics.
 
     Returns:
-        The system dynamics closed over the model, the data, and the additional kwargs.
+        The system dynamics closed over the additional kwargs to be used by
+        JaxSim integrators.
     """
 
-    # We allow to close `system_dynamics` over additional kwargs.
-    kwargs_closed = kwargs.copy()
-
-    # Create a local copy of model and data.
-    # The wrapped dynamics will hold a reference of this object.
-    model_closed = model.copy()
-    data_closed = data.copy().replace(
-        state=js.ode_data.ODEState.zero(model=model_closed, data=data)
-    )
-
+    # Close `system_dynamics` over additional kwargs.
+    # Similarly to what done in `jaxsim.api.model.step`, to be future-proof, we use the
+    # following logic to allow the caller to close over arguments having the same name
+    # of the ones used in the `wrap_system_dynamics_for_integration` function.
+    kwargs = kwargs.copy() if kwargs is not None else {}
+    colliding_system_dynamics_kwargs = kwargs.pop("system_dynamics_kwargs", {})
+    system_dynamics_kwargs = kwargs | colliding_system_dynamics_kwargs
+
+    # Remove `model` and `data` for backward compatibility.
+    # It's no longer necessary to close over them at this stage, as this is always
+    # done in `jaxsim.api.model.step`.
+    # We can remove the following lines in a few releases.
+    _ = system_dynamics_kwargs.pop("data", None)
+    _ = system_dynamics_kwargs.pop("model", None)
+
+    # Create the function with the signature expected by our generic integrators.
+    # Note that our system dynamics is time independent.
     def f(x: ODEState, t: Time, **kwargs_f) -> tuple[ODEState, dict[str, Any]]:
 
-        # Allow caller to override the closed data and model objects.
-        data_f = kwargs_f.pop("data", data_closed)
-        model_f = kwargs_f.pop("model", model_closed)
+        # Get the data and model objects from the kwargs.
+        data_f = kwargs_f.pop("data")
+        model_f = kwargs_f.pop("model")
 
         # Update the state and time stored inside data.
         with data_f.editable(validate=True) as data_rw:
@@ -69,7 +73,7 @@ def wrap_system_dynamics_for_integration(
         return system_dynamics(
             model=model_f,
             data=data_rw,
-            **(kwargs_closed | kwargs_f),
+            **(system_dynamics_kwargs | kwargs_f),
         )
 
     f: jaxsim.integrators.common.SystemDynamics[ODEState, ODEState]