imt-ring 1.6.37__tar.gz → 1.6.39__tar.gz

{imt_ring-1.6.37 → imt_ring-1.6.39}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: imt-ring
-Version: 1.6.37
+Version: 1.6.39
 Summary: RING: Recurrent Inertial Graph-based Estimator
 Author-email: Simon Bachhuber <simon.bachhuber@fau.de>
 Project-URL: Homepage, https://github.com/SimiPixel/ring

{imt_ring-1.6.37 → imt_ring-1.6.39}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "imt-ring"
-version = "1.6.37"
+version = "1.6.39"
 authors = [
   { name="Simon Bachhuber", email="simon.bachhuber@fau.de" },
 ]

{imt_ring-1.6.37 → imt_ring-1.6.39}/src/imt_ring.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: imt-ring
-Version: 1.6.37
+Version: 1.6.39
 Summary: RING: Recurrent Inertial Graph-based Estimator
 Author-email: Simon Bachhuber <simon.bachhuber@fau.de>
 Project-URL: Homepage, https://github.com/SimiPixel/ring

{imt_ring-1.6.37 → imt_ring-1.6.39}/src/ring/algorithms/custom_joints/suntay.py

@@ -184,7 +184,7 @@ def register_suntay(sconfig: SuntayConfig, name: str = "suntay"):
 
         suntay_link_name = _utils_find_suntay_joint(sys)
 
-        params = jax.tree_map(
+        params = jax.tree.map(
             lambda arr: arr[sys.idx_map("l")[suntay_link_name]],
             sys.links.joint_params[name],
         )

{imt_ring-1.6.37 → imt_ring-1.6.39}/src/ring/algorithms/dynamics.py

@@ -303,7 +303,33 @@ def step(
     taus: Optional[jax.Array] = None,
     n_substeps: int = 1,
 ) -> base.State:
-    "Steps the dynamics. Returns the state of next timestep."
+    """
+    Advances the system dynamics by a single timestep using semi-implicit Euler integration.
+
+    This function updates the system's state by integrating the equations of motion
+    over a timestep, potentially with multiple substeps for improved numerical stability.
+    The method ensures that the system's kinematics are updated before each integration step.
+
+    Args:
+        sys (base.System):
+            The system to simulate, containing link information, joint dynamics, and integration parameters.
+        state (base.State):
+            The current state of the system, including joint positions (`q`), velocities (`qd`), and transforms (`x`).
+        taus (Optional[jax.Array], optional):
+            The control torques applied to the system joints. If `None`, zero torques are applied.
+            Defaults to `None`.
+        n_substeps (int, optional):
+            The number of integration substeps per timestep to improve numerical accuracy.
+            Defaults to `1`.
+
+    Returns:
+        base.State:
+            The updated state of the system after integration.
+
+    Raises:
+        AssertionError: If the system's degrees of freedom (`q` and `qd`) do not match expectations.
+        AssertionError: If an unsupported integration method is specified in `sys.integration_method`.
+    """  # noqa: E501
     assert sys.q_size() == state.q.size
     if taus is None:
         taus = jnp.zeros_like(state.qd)

{imt_ring-1.6.37 → imt_ring-1.6.39}/src/ring/algorithms/generator/base.py

@@ -53,7 +53,85 @@ class RCMG:
         cor: bool = False,
         disable_tqdm: bool = False,
     ) -> None:
-        "Random Chain Motion Generator"
+        """
+        Initializes the Random Chain Motion Generator (RCMG).
+
+        The RCMG generates synthetic joint motion sequences for kinematic and dynamic
+        systems based on predefined motion configurations. It allows for system
+        randomization, augmentation with IMU and joint axis data, and optional
+        dynamic simulation.
+
+        Args:
+            sys (base.System | list[base.System]):
+                The system(s) for which motion should be generated.
+            config (jcalc.MotionConfig | list[jcalc.MotionConfig], optional):
+                Motion configuration(s) defining velocity limits, interpolation methods,
+                and range constraints. Defaults to `jcalc.MotionConfig()`.
+            setup_fn (Optional[types.SETUP_FN], optional):
+                A function to modify the system before motion generation. Defaults to `None`.
+            finalize_fn (Optional[types.FINALIZE_FN], optional):
+                A function to modify outputs after motion generation. Defaults to `None`.
+            add_X_imus (bool, optional):
+                Whether to add IMU sensor data to the output. Defaults to `False`.
+            add_X_imus_kwargs (dict, optional):
+                Additional keyword arguments for IMU data processing. Defaults to `{}`.
+            add_X_jointaxes (bool, optional):
+                Whether to add joint axis data to the output. Defaults to `False`.
+            add_X_jointaxes_kwargs (dict, optional):
+                Additional keyword arguments for joint axis data processing. Defaults to `{}`.
+            add_y_relpose (bool, optional):
+                Whether to add relative pose targets to the output. Defaults to `False`.
+            add_y_rootincl (bool, optional):
+                Whether to add root inclination targets to the output. Defaults to `False`.
+            add_y_rootincl_kwargs (dict, optional):
+                Additional keyword arguments for root inclination processing. Defaults to `{}`.
+            add_y_rootfull (bool, optional):
+                Whether to add full root state targets to the output. Defaults to `False`.
+            add_y_rootfull_kwargs (dict, optional):
+                Additional keyword arguments for full root state processing. Defaults to `{}`.
+            sys_ml (Optional[base.System], optional):
+                System that defines the graph and naming structure of the `X` and `y` outputs. Defaults to `None` which then uses the first provided system.
+            randomize_positions (bool, optional):
+                Whether to randomised positions based on `pos_min` and `pos_max`. Defaults to `False`.
+            randomize_motion_artifacts (bool, optional):
+                Whether to randomize the IMU motion artifact simulation. This randomises the spring stiffness and spring damping parameters of the passive free joint that is added between nonrigid and rigid IMU. Defaults to `False`.
+            randomize_joint_params (bool, optional):
+                Whether to randomize joint parameters by calling `JointModel.init_joint_params` before every sequence generation. Defaults to `False`.
+            randomize_hz (bool, optional):
+                Whether to randomize the sampling frequency of the generated data. Defaults to `False`.
+            randomize_hz_kwargs (dict, optional):
+                Additional keyword arguments for sampling frequency randomization. Defaults to `{}`.
+            imu_motion_artifacts (bool, optional):
+                Whether to simulate nonrigid IMU motion artifacts. Defaults to `False`.
+            imu_motion_artifacts_kwargs (dict, optional):
+                Additional keyword arguments for IMU motion artifact simulation. Defaults to `{}`.
+            dynamic_simulation (bool, optional):
+                Whether to use a physics-based simulation to generate motion instead of purely
+                kinematic methods. Defaults to `False`.
+            dynamic_simulation_kwargs (dict, optional):
+                Additional keyword arguments for dynamic simulation. Defaults to `{}`.
+            output_transform (Optional[Callable], optional):
+                A function to transform the generated output data. Defaults to `None`.
+            keep_output_extras (bool, optional):
+                Whether to keep additional output metadata. Defaults to `False`.
+            use_link_number_in_Xy (bool, optional):
+                Whether to replace joint names with numerical indices in the output. Defaults to `False`.
+            cor (bool, optional):
+                Whether to replace free joints with center-of-rotation (COR) 9D free joint. Defaults to `False`.
+            disable_tqdm (bool, optional):
+                Whether to disable progress bars during generation. Defaults to `False`.
+
+        Raises:
+            AssertionError: If any of the provided `MotionConfig` instances are infeasible.
+
+        Notes:
+            - This class supports batch generation, lazy and eager data loading, and
+              motion augmentation.
+            - If `randomize_hz=True`, the time step (`dt`) varies according to the specified
+              sampling rates.
+            - When `cor=True`, free joints are replaced with center-of-rotation models,
+              affecting joint motion behavior.
+        """  # noqa: E501
 
         # add some default values
         randomize_hz_kwargs_defaults = dict(add_dt=True)
@@ -139,6 +217,7 @@ class RCMG:
     def to_lazy_gen(
         self, sizes: int | list[int] = 1, jit: bool = True
     ) -> types.BatchedGenerator:
+        "Returns a generator `X, y = gen(key)` that lazily generates batched sequences."
         return batch.generators_lazy(self.gens, self._compute_repeats(sizes), jit)
 
     @staticmethod
@@ -201,7 +280,7 @@ class RCMG:
         ),
         verbose: bool = True,
     ):
-
+        "Stores unbatched sequences as numpy arrays into folder."
         i = 0
 
         def callback(data: list[PyTree[np.ndarray]]) -> None:
@@ -237,6 +316,7 @@ class RCMG:
         shuffle: bool = True,
         transform=None,
     ) -> types.BatchedGenerator:
+        "Returns a generator `X, y = gen(key)` that returns precomputed batched sequences."  # noqa: E501
         data = self.to_list(sizes, seed)
         assert len(data) >= batchsize
         return self.eager_gen_from_list(data, batchsize, shuffle, transform)

{imt_ring-1.6.37 → imt_ring-1.6.39}/src/ring/algorithms/generator/batch.py

@@ -80,11 +80,11 @@ def generators_eager(
             # converts also to numpy; but with np.array.flags.writeable = False
             sample = jax.device_get(sample)
             # this then sets this flag to True
-            sample = jax.tree_map(np.array, sample)
+            sample = jax.tree.map(np.array, sample)
 
             sample_flat, _ = jax.tree_util.tree_flatten(sample)
             size = 1 if len(sample_flat) == 0 else sample_flat[0].shape[0]
-            callback([jax.tree_map(lambda a: a[i].copy(), sample) for i in range(size)])
+            callback([jax.tree.map(lambda a: a[i].copy(), sample) for i in range(size)])
 
             # cleanup
             del sample, sample_flat

{imt_ring-1.6.37 → imt_ring-1.6.39}/src/ring/algorithms/generator/finalize_fns.py

@@ -311,7 +311,7 @@ def _expand_then_flatten(Xy):
 
     X, y = _flatten(X), _flatten(y)
     if not batched:
-        X, y = jax.tree_map(lambda arr: arr[0], (X, y))
+        X, y = jax.tree.map(lambda arr: arr[0], (X, y))
     return X, y
 
 

{imt_ring-1.6.37 → imt_ring-1.6.39}/src/ring/algorithms/generator/pd_control.py

@@ -86,7 +86,7 @@ def _pd_control(P: jax.Array, D: Optional[jax.Array] = None):
         controller_state: PDControllerState, sys: base.System, state: base.State
     ) -> jax.Array:
         taus = jnp.zeros((sys.qd_size()))
-        q_ref, qd_ref = jax.tree_map(
+        q_ref, qd_ref = jax.tree.map(
             lambda arr: jax.lax.dynamic_index_in_dim(
                 arr, controller_state.i, keepdims=False
             ),

{imt_ring-1.6.37 → imt_ring-1.6.39}/src/ring/algorithms/jcalc.py

@@ -19,6 +19,87 @@ from ring.algorithms._random import TimeDependentFloat
 
 @dataclass
 class MotionConfig:
+    """
+    Configuration for joint motion generation in kinematic and dynamic simulations.
+
+    This class defines the constraints and parameters for generating random joint motions,
+    including angular and positional velocity limits, interpolation methods, and range
+    restrictions for various joint types.
+
+    Attributes:
+        T (float): Total duration of the motion sequence (in seconds).
+        t_min (float): Minimum time interval between two generated joint states.
+        t_max (float | TimeDependentFloat): Maximum time interval between two generated joint states.
+
+        dang_min (float | TimeDependentFloat): Minimum angular velocity (rad/s).
+        dang_max (float | TimeDependentFloat): Maximum angular velocity (rad/s).
+        dang_min_free_spherical (float | TimeDependentFloat): Minimum angular velocity for free and spherical joints.
+        dang_max_free_spherical (float | TimeDependentFloat): Maximum angular velocity for free and spherical joints.
+
+        delta_ang_min (float | TimeDependentFloat): Minimum allowed change in joint angle (radians).
+        delta_ang_max (float | TimeDependentFloat): Maximum allowed change in joint angle (radians).
+        delta_ang_min_free_spherical (float | TimeDependentFloat): Minimum allowed change in angle for free/spherical joints.
+        delta_ang_max_free_spherical (float | TimeDependentFloat): Maximum allowed change in angle for free/spherical joints.
+
+        dpos_min (float | TimeDependentFloat): Minimum translational velocity.
+        dpos_max (float | TimeDependentFloat): Maximum translational velocity.
+        pos_min (float | TimeDependentFloat): Minimum position constraint.
+        pos_max (float | TimeDependentFloat): Maximum position constraint.
+
+        pos_min_p3d_x (float | TimeDependentFloat): Minimum position in x-direction for P3D joints.
+        pos_max_p3d_x (float | TimeDependentFloat): Maximum position in x-direction for P3D joints.
+        pos_min_p3d_y (float | TimeDependentFloat): Minimum position in y-direction for P3D joints.
+        pos_max_p3d_y (float | TimeDependentFloat): Maximum position in y-direction for P3D joints.
+        pos_min_p3d_z (float | TimeDependentFloat): Minimum position in z-direction for P3D joints.
+        pos_max_p3d_z (float | TimeDependentFloat): Maximum position in z-direction for P3D joints.
+
+        cdf_bins_min (int): Minimum number of bins for cumulative distribution function (CDF)-based random sampling.
+        cdf_bins_max (Optional[int]): Maximum number of bins for CDF-based sampling.
+
+        randomized_interpolation_angle (bool): Whether to use randomized interpolation for angular motion.
+        randomized_interpolation_position (bool): Whether to use randomized interpolation for positional motion.
+        interpolation_method (str): Interpolation method to be used (default: "cosine").
+
+        range_of_motion_hinge (bool): Whether to enforce range-of-motion constraints on hinge joints.
+        range_of_motion_hinge_method (str): Method used for range-of-motion enforcement (e.g., "uniform", "sigmoid").
+
+        rom_halfsize (float | TimeDependentFloat): Half-size of the range of motion restriction.
+
+        ang0_min (float): Minimum initial joint angle.
+        ang0_max (float): Maximum initial joint angle.
+        pos0_min (float): Minimum initial joint position.
+        pos0_max (float): Maximum initial joint position.
+
+        cor_t_min (float): Minimum time step for center-of-rotation (COR) joints.
+        cor_t_max (float | TimeDependentFloat): Maximum time step for COR joints.
+        cor_dpos_min (float | TimeDependentFloat): Minimum velocity for COR translation.
+        cor_dpos_max (float | TimeDependentFloat): Maximum velocity for COR translation.
+        cor_pos_min (float | TimeDependentFloat): Minimum position for COR translation.
+        cor_pos_max (float | TimeDependentFloat): Maximum position for COR translation.
+        cor_pos0_min (float): Initial minimum position for COR translation.
+        cor_pos0_max (float): Initial maximum position for COR translation.
+
+        joint_type_specific_overwrites (dict[str, dict[str, Any]]):
+            A dictionary mapping joint types to specific motion configuration overrides.
+
+    Methods:
+        is_feasible:
+            Checks if the motion configuration satisfies all constraints.
+
+        to_nomotion_config:
+            Returns a new `MotionConfig` where all velocities and angle changes are set to zero.
+
+        overwrite_for_joint_type:
+            Applies specific configuration changes for a given joint type.
+            Note: These changes affect all instances of `MotionConfig` for this joint type.
+
+        overwrite_for_subsystem:
+            Modifies the motion configuration for all joints in a subsystem rooted at `link_name`.
+
+        from_register:
+            Retrieves a predefined `MotionConfig` from the global registry.
+    """  # noqa: E501
+
     T: float = 60.0  # length of random motion
     t_min: float = 0.05  # min time between two generated angles
     t_max: float | TimeDependentFloat = 0.30  # max time ..
@@ -412,6 +493,30 @@ def _find_interval(t: jax.Array, boundaries: jax.Array):
 def join_motionconfigs(
     configs: list[MotionConfig], boundaries: list[float]
 ) -> MotionConfig:
+    """
+    Joins multiple `MotionConfig` objects in time, transitioning between them at specified boundaries.
+
+    This function takes a list of `MotionConfig` instances and a corresponding list of boundary times,
+    and constructs a new `MotionConfig` that varies in time according to the provided segments.
+
+    Args:
+        configs (list[MotionConfig]): A list of `MotionConfig` objects to be joined.
+        boundaries (list[float]): A list of time values where transitions between `configs` occur.
+            Must have one element less than `configs`, as each boundary defines the transition point
+            between two consecutive configurations.
+
+    Returns:
+        MotionConfig: A new `MotionConfig` object where time-dependent fields transition based on the
+        specified boundaries.
+
+    Raises:
+        AssertionError: If the number of boundaries does not match `len(configs) - 1`.
+        AssertionError: If time-independent fields have differing values across `configs`.
+
+    Notes:
+        - Only fields that are time-dependent (`float | TimeDependentFloat`) will change over time.
+        - Time-independent fields must be the same in all `configs`, or an error is raised.
+    """  # noqa: E501
     # to avoid a circular import due to `ring.utils.randomize_sys` importing `jcalc`
     from ring.utils import tree_equal
 
@@ -517,6 +622,55 @@ INV_KIN = Callable[[base.Transform, tree_utils.PyTree], jax.Array]
 
 @dataclass
 class JointModel:
+    """
+    Represents the kinematic and dynamic properties of a joint type.
+
+    A `JointModel` defines the mathematical functions required to compute joint
+    transformations, motion, control terms, and inverse kinematics. It is used to
+    describe the behavior of various joint types, including revolute, prismatic,
+    spherical, and free joints.
+
+    Attributes:
+        transform (Callable[[jax.Array, jax.Array], base.Transform]):
+            Computes the transformation (position and orientation) of the joint
+            given the joint state `q` and joint parameters.
+
+        motion (list[base.Motion | Callable[[jax.Array], base.Motion]]):
+            Defines the joint motion model. It can be a list of `Motion` objects
+            or callables that return `Motion` based on joint parameters.
+
+        rcmg_draw_fn (Optional[DRAW_FN]):
+            Function used to generate a reference motion trajectory for the joint
+            using Randomized Control Motion Generation (RCMG).
+
+        p_control_term (Optional[P_CONTROL_TERM]):
+            Function that computes the proportional control term for the joint.
+
+        qd_from_q (Optional[QD_FROM_Q]):
+            Function to compute joint velocity (`qd`) from joint positions (`q`).
+
+        coordinate_vector_to_q (Optional[COORDINATE_VECTOR_TO_Q]):
+            Function that maps a coordinate vector to a valid joint state `q`,
+            ensuring constraints (e.g., wrapping angles or normalizing quaternions).
+
+        inv_kin (Optional[INV_KIN]):
+            Function that computes the inverse kinematics for the joint, mapping
+            a desired transform to joint coordinates `q`.
+
+        init_joint_params (Optional[INIT_JOINT_PARAMS]):
+            Function that initializes joint-specific parameters.
+
+        utilities (Optional[dict[str, Any]]):
+            Additional utility functions or metadata related to the joint model.
+
+    Notes:
+        - The `transform` function is essential for computing the joint's spatial
+          transformation based on its generalized coordinates.
+        - The `motion` attribute describes how forces and torques affect the joint.
+        - The `rcmg_draw_fn` is used for RCMG motion generation.
+        - The `coordinate_vector_to_q` is critical for maintaining valid joint states.
+    """  # noqa: E501
+
     # (q, params) -> Transform
     transform: Callable[[jax.Array, jax.Array], base.Transform]
     # len(motion) == len(qd)
@@ -1079,6 +1233,50 @@ def register_new_joint_type(
     qd_width: Optional[int] = None,
     overwrite: bool = False,
 ):
+    """
+    Registers a new joint type with its corresponding `JointModel` and kinematic properties.
+
+    This function allows the addition of custom joint types to the system by associating
+    them with a `JointModel`, specifying their state and velocity dimensions, and optionally
+    overwriting existing joint definitions.
+
+    Args:
+        joint_type (str):
+            Name of the new joint type to register.
+        joint_model (JointModel):
+            The `JointModel` instance defining the kinematic and dynamic properties of the joint.
+        q_width (int):
+            Number of generalized coordinates (degrees of freedom) required to represent the joint.
+        qd_width (Optional[int], default=None):
+            Number of velocity coordinates associated with the joint. Defaults to `q_width`.
+        overwrite (bool, default=False):
+            If `True`, allows overwriting an existing joint type. Otherwise, raises an error if
+            the joint type already exists.
+
+    Raises:
+        AssertionError:
+            - If `joint_type` is `"default"` (reserved name).
+            - If `joint_type` already exists and `overwrite=False`.
+            - If `qd_width` is not provided and does not default to `q_width`.
+            - If `joint_model.motion` length does not match `qd_width`.
+
+    Notes:
+        - The function updates global dictionaries that store joint properties, including:
+          - `_joint_types`: Maps joint type names to `JointModel` instances.
+          - `base.Q_WIDTHS`: Stores the number of state coordinates for each joint type.
+          - `base.QD_WIDTHS`: Stores the number of velocity coordinates for each joint type.
+        - If `overwrite=True`, existing entries are removed before adding the new joint type.
+        - Ensures consistency between motion definitions and velocity coordinate dimensions.
+
+    Example:
+        ```python
+        new_joint = JointModel(
+            transform=my_transform_fn,
+            motion=[base.Motion.create(ang=jnp.array([1, 0, 0]))],
+        )
+        register_new_joint_type("custom_hinge", new_joint, q_width=1)
+        ```
+    """  # noqa: E501
     # this name is used
     assert joint_type != "default", "Please use another name."
 

{imt_ring-1.6.37 → imt_ring-1.6.39}/src/ring/algorithms/kinematics.py

@@ -4,6 +4,7 @@ import jax
 import jax.numpy as jnp
 import jaxopt
 from jaxopt._src.base import Solver
+
 from ring import algebra
 from ring import base
 from ring import maths
@@ -171,7 +172,7 @@ def inverse_kinematics_endeffector(
 
             # find result of best q0 initial value
             best_q_index = jnp.argmin(values)
-            best_q, best_q_value = jax.tree_map(
+            best_q, best_q_value = jax.tree.map(
                 lambda arr: jax.lax.dynamic_index_in_dim(
                     arr, best_q_index, keepdims=False
                 ),

{imt_ring-1.6.37 → imt_ring-1.6.39}/src/ring/algorithms/sensors.py

@@ -244,7 +244,7 @@ def imu(
         measurements["mag"] = magnetometer(xs.rot, magvec)
 
     if smoothen_degree is not None:
-        measurements = jax.tree_map(
+        measurements = jax.tree.map(
             lambda arr: _moving_average(arr, smoothen_degree),
             measurements,
         )
@@ -257,7 +257,7 @@ def imu(
             delay = half_window
 
     if delay is not None and delay > 0:
-        measurements = jax.tree_map(
+        measurements = jax.tree.map(
             lambda arr: (jnp.pad(arr, ((delay, 0), (0, 0)))[:-delay]), measurements
         )
 
@@ -473,7 +473,7 @@ def _joint_axes_from_sys(sys: base.Transform, N: int) -> dict:
         X[name] = {"joint_axes": joint_axes}
 
     sys.scan(f, "lll", sys.link_names, sys.link_types, sys.links)
-    X = jax.tree_map(lambda arr: jnp.repeat(arr[None], N, axis=0), X)
+    X = jax.tree.map(lambda arr: jnp.repeat(arr[None], N, axis=0), X)
     return X
 
 
@@ -498,12 +498,12 @@ _quasi_physical_sys_str = r"""
 <x_xy>
     <options gravity="0 0 0"/>
     <worldbody>
-        <body name="IMU" joint="p3d" damping="0.1 0.1 0.1" spring_stiff="3 3 3">
-            <geom type="box" mass="0.002" dim="0.01 0.01 0.01"/>
+        <body name="IMU" joint="free" damping="1 1 1 10 10 10" spring_stiff="20 20 20 500 500 500">
+            <geom type="box" mass="1" dim="0.01 0.01 0.01"/>
         </body>
     </worldbody>
 </x_xy>
-"""
+"""  # noqa: E501
 
 
 def _quasi_physical_simulation_beautiful(
@@ -512,12 +512,14 @@ def _quasi_physical_simulation_beautiful(
     sys = io.load_sys_from_str(_quasi_physical_sys_str).replace(dt=dt)
 
     def step_dynamics(state: base.State, x):
-        state = algorithms.step(sys.replace(link_spring_zeropoint=x.pos), state)
+        state = algorithms.step(
+            sys.replace(link_spring_zeropoint=jnp.concatenate((x.rot, x.pos))), state
+        )
         return state, state.q
 
-    state = base.State.create(sys, q=xs.pos[0])
-    _, pos = jax.lax.scan(step_dynamics, state, xs)
-    return xs.replace(pos=pos)
+    state = base.State.create(sys, q=jnp.concatenate((xs.rot[0], xs.pos[0])))
+    _, qs = jax.lax.scan(step_dynamics, state, xs)
+    return xs.replace(rot=qs[:, :4], pos=qs[:, 4:])
 
 
 _constants = {