imt-ring 1.6.38__py3-none-any.whl → 1.6.39__py3-none-any.whl

{imt_ring-1.6.38.dist-info → imt_ring-1.6.39.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: imt-ring
-Version: 1.6.38
+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.38.dist-info → imt_ring-1.6.39.dist-info}/RECORD

@@ -1,12 +1,12 @@
 ring/__init__.py,sha256=H1Rd2uXVkux4Z792XyHIkQ8OpDSZBiPqFwyAFDWDU3E,5260
 ring/algebra.py,sha256=F0GwbP8LQP5qGVkoMUYJmkp9Hn2nKAVIkCVYDEjNjGU,3128
-ring/base.py,sha256=4Yxk6jk-B4UUm_6YYshxmHSHqOg0mhTOxtZP5fFS8nw,35373
+ring/base.py,sha256=zromjIuMpNBoyiwHa9OCyZvAz7jHjXHZIdRt8fN8PoA,50481
 ring/maths.py,sha256=R22SNQutkf9v7Hp9klo0wvJVIyBQz0O8_5oJaDQcFis,12652
 ring/spatial.py,sha256=nmZ-UhRanhyM34bez8uCS4wMwaKqLkuEbgKGP5XNH60,2351
 ring/algorithms/__init__.py,sha256=IiK9EN5Xgs3dB075-A-H-Yad0Z7vzvKIJF2g6X_-C_8,1224
 ring/algorithms/_random.py,sha256=UMyv-VPZLcErrKqs0XB83QJjs8GrmoNsv-zRSxGXvnI,14490
-ring/algorithms/dynamics.py,sha256=dpe-F3Yq4sY2dY6DQW3v7TnPLRdxdkePtdbGPQIrijg,10997
-ring/algorithms/jcalc.py,sha256=QafnCKa1mjEl7nL_KuadPJB5ebW31NKnkdcKn2YtSsM,36171
+ring/algorithms/dynamics.py,sha256=NFOZawjrFoS5RgiWOpG1pQCC8l7RBOEZIi9ok6gvf9U,12268
+ring/algorithms/jcalc.py,sha256=l6BXOmXwrZ_AKKRm4gEHq_k2LSUQ4wd--1qL1qNTcKk,46794
 ring/algorithms/kinematics.py,sha256=IXeTQ-afzeEzLVmQVQ1oTXJxz_lTwvaWlgHeJkhO_8o,7423
 ring/algorithms/sensors.py,sha256=v_TZMyWjffDpPwoyS1fy8X-1i9y1vDf6mk1EmGS2ztc,18251
 ring/algorithms/custom_joints/__init__.py,sha256=3pQ-Is_HBTQDkzESCNg9VfoP8wvseWmooryG8ERnu_A,366
@@ -15,7 +15,7 @@ ring/algorithms/custom_joints/rr_joint.py,sha256=jnRtjtOCALMaq2_0bcu2d7qgfQ6etXp
 ring/algorithms/custom_joints/rsaddle_joint.py,sha256=QoMo6NXdYgA9JygSzBvr0eCdd3qKhUgCrGPNO2Qdxko,1200
 ring/algorithms/custom_joints/suntay.py,sha256=TZG307NqdMiXnNY63xEx8AkAjbQBQ4eO6DQ7R4j4D08,16726
 ring/algorithms/generator/__init__.py,sha256=bF-CW3x2x-o6KWESKy-DuxzZPh3UNSjJb_MaAcSHGsQ,277
-ring/algorithms/generator/base.py,sha256=jGQocoNZ5tkiMazBDCv-jD6FNYwebqn0_RgVFse49pg,16890
+ring/algorithms/generator/base.py,sha256=klWYt6TlMluLu0ihGzmmPXBm47DOTpjXJylZVNXHVEk,22419
 ring/algorithms/generator/batch.py,sha256=xp1X8oYtwI6l2cH4GRu9zw-P8dnh-X1FWTSyixEfgr8,2652
 ring/algorithms/generator/finalize_fns.py,sha256=ty1NaU-Mghx1RL-voivDjS0TWSKNtjTmbdmBnShhn7k,10398
 ring/algorithms/generator/motion_artifacts.py,sha256=2VJbldVDbI3PSyboshIbtYvSAKzBBwGV7cQfYjqvluM,9167
@@ -52,7 +52,7 @@ ring/io/xml/test_from_xml.py,sha256=bckVrVVmEhCwujd_OF9FGYnX3zU3BgztpqGxxmd0htM,
 ring/io/xml/test_to_xml.py,sha256=NGn4VSiFdwhYN5YTBduWMiY9B5dwtxZhCQAR_PXeqKU,946
 ring/io/xml/to_xml.py,sha256=Wo4iySLw9nM-iVW42AGvMRqjtU2qRc2FD_Zlc7w1IrE,3438
 ring/ml/__init__.py,sha256=nbh48gaswWeY4S4vT1sply_3ROj2DQ7agjoLR4Ho3T8,1517
-ring/ml/base.py,sha256=eEpLuXlhFoJ-cpXoSGjctLaYduQhnSVpbv-FEYftNRs,9972
+ring/ml/base.py,sha256=phBUfTpP1Mqt8lvtilSavT7ypkaeaF1oh7nCMjV0dqg,10046
 ring/ml/callbacks.py,sha256=oCPXl4_Zcw3g0KRgyyUDmdiGxV0phnDVc_t8rEG4Lls,13737
 ring/ml/ml_utils.py,sha256=hu189AnHcmkhkpEPZZ19O0gWz3T-YKpWQW9buqDTMow,10915
 ring/ml/optimizer.py,sha256=TZF0_LmnewzmGVso-zIQJtpWguUW0fW3HeRpIdG_qoI,4763
@@ -86,7 +86,7 @@ ring/utils/randomize_sys.py,sha256=G_vBIo0OwQkXL2u0djwbaoaeb02C4LQCTNNloOYIU2M,3
 ring/utils/utils.py,sha256=gKwOXLxWraeZfX6EbBcg3hkq30DcXN0mcRUeOSTNiMo,7336
 ring/utils/register_gym_envs/__init__.py,sha256=PtPIRBQJ16339xZ9G9VpvqrvcGbQ_Pk_SUz4tQPa9nQ,94
 ring/utils/register_gym_envs/saddle.py,sha256=tA5CyW_akSXyDm0xJ83CtOrUMVElH0f9vZtEDDJQalI,4422
-imt_ring-1.6.38.dist-info/METADATA,sha256=9rN1VzsIlGU8eyABz9-pTxj0OTCFOZRilEEzkB4gyvg,4251
-imt_ring-1.6.38.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-imt_ring-1.6.38.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
-imt_ring-1.6.38.dist-info/RECORD,,
+imt_ring-1.6.39.dist-info/METADATA,sha256=v0rBTnCQP-SWJU153byfX31HUCcrWHhCg_EmecLbLf4,4251
+imt_ring-1.6.39.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+imt_ring-1.6.39.dist-info/top_level.txt,sha256=EiT790-lAyi8iwTzJArH3f2k77rwhDn00q-4PlmvDQo,5
+imt_ring-1.6.39.dist-info/RECORD,,

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)

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)

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."
 

ring/base.py

@@ -112,17 +112,71 @@ class _Base:
 
 @struct.dataclass
 class Transform(_Base):
-    """Represents the Transformation from Plücker A to Plücker B,
-    where B is located relative to A at `pos` in frame A and `rot` is the
-    relative quaternion from A to B.
-    Create using `Transform.create(pos=..., rot=...)
     """
+    Represents a spatial transformation between two coordinate frames using Plücker coordinates.
+
+    The `Transform` class defines the relative position and orientation of one frame (`B`)
+    with respect to another frame (`A`). The position (`pos`) is given in the coordinate frame
+    of `A`, and the rotation (`rot`) is expressed as a unit quaternion representing the relative
+    rotation from frame `A` to frame `B`.
+
+    Attributes:
+        pos (Vector):
+            The translation vector (position of `B` relative to `A`) in the coordinate frame of `A`.
+            Shape: `(..., 3)`, where `...` represents optional batch dimensions.
+        rot (Quaternion):
+            The unit quaternion representing the orientation of `B` relative to `A`.
+            Shape: `(..., 4)`, where `...` represents optional batch dimensions.
+
+    Methods:
+        create(pos: Optional[Vector] = None, rot: Optional[Quaternion] = None) -> Transform:
+            Creates a `Transform` instance with optional position and rotation.
+
+        zero(shape: Sequence[int] = ()) -> Transform:
+            Returns a zero transform with a given batch shape.
+
+        as_matrix() -> jax.Array:
+            Returns the 4x4 homogeneous transformation matrix representation of this transform.
+
+    Usage:
+        >>> pos = jnp.array([1.0, 2.0, 3.0])
+        >>> rot = jnp.array([1.0, 0.0, 0.0, 0.0])  # Identity quaternion
+        >>> T = Transform.create(pos, rot)
+        >>> print(T.pos)  # Output: [1. 2. 3.]
+        >>> print(T.rot)  # Output: [1. 0. 0. 0.]
+        >>> print(T.as_matrix())  # 4x4 transformation matrix
+    """  # noqa: E501
 
     pos: Vector
     rot: Quaternion
 
     @classmethod
     def create(cls, pos=None, rot=None):
+        """
+        Creates a `Transform` instance with the specified position and rotation.
+
+        At least one of `pos` or `rot` must be provided. If only `pos` is given, the rotation
+        defaults to the identity quaternion `[1, 0, 0, 0]`. If only `rot` is given, the position
+        defaults to `[0, 0, 0]`.
+
+        Args:
+            pos (Optional[Vector], default=None):
+                The position of frame `B` relative to frame `A`, expressed in frame `A` coordinates.
+                If `None`, defaults to a zero vector of shape `(3,)`.
+            rot (Optional[Quaternion], default=None):
+                The unit quaternion representing the orientation of `B` relative to `A`.
+                If `None`, defaults to the identity quaternion `(1, 0, 0, 0)`.
+
+        Returns:
+            Transform: A new `Transform` instance with the specified position and rotation.
+
+        Example:
+            >>> pos = jnp.array([1.0, 2.0, 3.0])
+            >>> rot = jnp.array([1.0, 0.0, 0.0, 0.0])  # Identity quaternion
+            >>> T = Transform.create(pos, rot)
+            >>> print(T.pos)  # Output: [1. 2. 3.]
+            >>> print(T.rot)  # Output: [1. 0. 0. 0.]
+        """  # noqa: E501
         assert not (pos is None and rot is None), "One must be given."
         shape_rot = rot.shape[:-1] if rot is not None else ()
         shape_pos = pos.shape[:-1] if pos is not None else ()
@@ -139,12 +193,49 @@ class Transform(_Base):
 
     @classmethod
     def zero(cls, shape=()) -> "Transform":
-        """Returns a zero transform with a batch shape."""
+        """
+        Returns a zero transform with a given batch shape.
+
+        This creates a transform with position `(0, 0, 0)` and an identity quaternion `(1, 0, 0, 0)`,
+        which represents no translation or rotation.
+
+        Args:
+            shape (Sequence[int], default=()):
+                The batch shape for the transform. Defaults to a scalar transform.
+
+        Returns:
+            Transform: A zero transform with the specified batch shape.
+
+        Example:
+            >>> T = Transform.zero()
+            >>> print(T.pos)  # Output: [0. 0. 0.]
+            >>> print(T.rot)  # Output: [1. 0. 0. 0.]
+        """  # noqa: E501
         pos = jnp.zeros(shape + (3,))
         rot = jnp.tile(jnp.array([1.0, 0.0, 0.0, 0.0]), shape + (1,))
         return Transform(pos, rot)
 
     def as_matrix(self) -> jax.Array:
+        """
+        Returns the 4x4 homogeneous transformation matrix representation of this transform.
+
+        The homogeneous transformation matrix is defined as:
+
+        ```
+        [ R  t ]
+        [ 0  1 ]
+        ```
+
+        where `R` is the 3x3 rotation matrix converted from the quaternion and `t` is the
+        3x1 position vector.
+
+        Returns:
+            jax.Array: A `(4, 4)` homogeneous transformation matrix.
+
+        Example:
+            >>> T = Transform.create(jnp.array([1.0, 2.0, 3.0]), jnp.array([1.0, 0.0, 0.0, 0.0]))
+            >>> print(T.as_matrix())  # Output: 4x4 matrix
+        """  # noqa: E501
         E = maths.quat_to_3x3(self.rot)
         return spatial.quadrants(aa=E, bb=E) @ spatial.xlt(self.pos)
 
@@ -402,7 +493,175 @@ QD_WIDTHS = {
 
 @struct.dataclass
 class System(_Base):
-    "System object. Create using `System.create(path_xml)`"
+    """
+    Represents a robotic system consisting of interconnected links and joints. Create it using `System.create(...)`
+
+    The `System` class models the kinematic and dynamic properties of a multibody
+    system, providing methods for state representation, transformations, joint
+    configuration management, and rendering. It supports both minimal and maximal
+    coordinate representations and can be parsed from or saved to XML files.
+
+    Attributes:
+        link_parents (list[int]):
+            A list specifying the parent index for each link. The root link has a parent index of `-1`.
+        links (Link):
+            A data structure containing information about all links in the system.
+        link_types (list[str]):
+            A list specifying the joint type for each link (e.g., "free", "hinge", "prismatic").
+        link_damping (jax.Array):
+            Joint damping coefficients for each link.
+        link_armature (jax.Array):
+            Armature inertia values for each joint.
+        link_spring_stiffness (jax.Array):
+            Stiffness values for joint springs.
+        link_spring_zeropoint (jax.Array):
+            Rest position of joint springs.
+        dt (float):
+            Simulation time step size.
+        geoms (list[Geometry]):
+            List of geometries associated with the system.
+        gravity (jax.Array):
+            Gravity vector applied to the system (default: `[0, 0, -9.81]`).
+        integration_method (str):
+            Integration method for simulation (default: "semi_implicit_euler").
+        mass_mat_iters (int):
+            Number of iterations for mass matrix calculations.
+        link_names (list[str]):
+            Names of the links in the system.
+        model_name (Optional[str]):
+            Name of the system model (if available).
+        omc (list[MaxCoordOMC | None]):
+            List of optional Maximal Coordinate representations.
+
+    Methods:
+        num_links() -> int:
+            Returns the number of links in the system.
+
+        q_size() -> int:
+            Returns the total number of generalized coordinates (`q`) in the system.
+
+        qd_size() -> int:
+            Returns the total number of generalized velocities (`qd`) in the system.
+
+        name_to_idx(name: str) -> int:
+            Returns the index of a link given its name.
+
+        idx_to_name(idx: int, allow_world: bool = False) -> str:
+            Returns the name of a link given its index. If `allow_world` is `True`,
+            returns `"world"` for index `-1`.
+
+        idx_map(type: str) -> dict:
+            Returns a dictionary mapping link names to their indices for a specified type
+            (`"l"`, `"q"`, or `"d"`).
+
+        parent_name(name: str) -> str:
+            Returns the name of the parent link for a given link.
+
+        change_model_name(new_name: Optional[str] = None, prefix: Optional[str] = None, suffix: Optional[str] = None) -> "System":
+            Changes the name of the system model.
+
+        change_link_name(old_name: str, new_name: str) -> "System":
+            Renames a specific link.
+
+        add_prefix_suffix(prefix: Optional[str] = None, suffix: Optional[str] = None) -> "System":
+            Adds both a prefix and suffix to all link names.
+
+        freeze(name: str | list[str]) -> "System":
+            Freezes the specified link(s), making them immovable.
+
+        unfreeze(name: str, new_joint_type: str) -> "System":
+            Unfreezes a frozen link and assigns it a new joint type.
+
+        change_joint_type(name: str, new_joint_type: str, **kwargs) -> "System":
+            Changes the joint type of a specified link.
+
+        joint_type_simplification(typ: str) -> str:
+            Returns a simplified representation of the given joint type.
+
+        joint_type_is_free_or_cor(typ: str) -> bool:
+            Checks if a joint type is either "free" or "cor".
+
+        joint_type_is_spherical(typ: str) -> bool:
+            Checks if a joint type is "spherical".
+
+        joint_type_is_free_or_cor_or_spherical(typ: str) -> bool:
+            Checks if a joint type is "free", "cor", or "spherical".
+
+        findall_imus(names: bool = True) -> list[str] | list[int]:
+            Finds all IMU sensors in the system.
+
+        findall_segments(names: bool = True) -> list[str] | list[int]:
+            Finds all non-IMU segments in the system.
+
+        findall_bodies_to_world(names: bool = False) -> list[int] | list[str]:
+            Returns all bodies directly connected to the world.
+
+        find_body_to_world(name: bool = False) -> int | str:
+            Returns the root body connected to the world.
+
+        findall_bodies_with_jointtype(typ: str, names: bool = False) -> list[int] | list[str]:
+            Returns all bodies with the specified joint type.
+
+        children(name: str, names: bool = False) -> list[int] | list[str]:
+            Returns the direct children of a given body.
+
+        findall_bodies_subsystem(name: str, names: bool = False) -> list[int] | list[str]:
+            Finds all bodies in the subsystem rooted at a given link.
+
+        scan(f: Callable, in_types: str, *args, reverse: bool = False):
+            Iterates over system elements while applying a function.
+
+        parse() -> "System":
+            Parses the system, performing consistency checks and computing spatial inertia tensors.
+
+        render(xs: Optional[Transform | list[Transform]] = None, **kwargs) -> list[np.ndarray]:
+            Renders frames of the system using maximal coordinates.
+
+        render_prediction(xs: Transform | list[Transform], yhat: dict | jax.Array | np.ndarray, **kwargs):
+            Renders a predicted state transformation.
+
+        delete_system(link_name: str | list[str], strict: bool = True):
+            Removes a subsystem from the system.
+
+        make_sys_noimu(imu_link_names: Optional[list[str]] = None):
+            Returns a version of the system without IMU sensors.
+
+        inject_system(other_system: "System", at_body: Optional[str] = None):
+            Merges another system into this one.
+
+        morph_system(new_parents: Optional[list[int | str]] = None, new_anchor: Optional[int | str] = None):
+            Reorders the system’s link hierarchy.
+
+        from_xml(path: str, seed: int = 1) -> "System":
+            Loads a system from an XML file.
+
+        from_str(xml: str, seed: int = 1) -> "System":
+            Loads a system from an XML string.
+
+        to_str(warn: bool = True) -> str:
+            Serializes the system to an XML string.
+
+        to_xml(path: str) -> None:
+            Saves the system as an XML file.
+
+        create(path_or_str: str, seed: int = 1) -> "System":
+            Creates a `System` instance from an XML file or string.
+
+        coordinate_vector_to_q(q: jax.Array, custom_joints: dict[str, Callable] = {}) -> jax.Array:
+            Converts a coordinate vector to minimal coordinates (`q`), applying
+            constraints such as quaternion normalization.
+
+    Raises:
+        AssertionError: If the system structure is invalid (e.g., duplicate link names, incorrect parent-child relationships).
+        InvalidSystemError: If an operation results in an inconsistent system state.
+
+    Notes:
+        - The system must be parsed before use to ensure consistency.
+        - The system supports batch operations using JAX for efficient computations.
+        - Joint types include revolute ("rx", "ry", "rz"), prismatic ("px", "py", "pz"), spherical, free, and more.
+        - Inertial properties of links are computed automatically from associated geometries.
+    """  # noqa: E501
+
     link_parents: list[int] = struct.field(False)
     links: Link
     link_types: list[str] = struct.field(False)
@@ -429,25 +688,32 @@ class System(_Base):
     omc: list[MaxCoordOMC | None] = struct.field(True, default_factory=lambda: [])
 
     def num_links(self) -> int:
+        "Returns the number of links in the system."
         return len(self.link_parents)
 
     def q_size(self) -> int:
+        "Returns the total number of generalized coordinates (`q`) in the system."
         return sum([Q_WIDTHS[typ] for typ in self.link_types])
 
     def qd_size(self) -> int:
+        "Returns the total number of generalized velocities (`qd`) in the system."
         return sum([QD_WIDTHS[typ] for typ in self.link_types])
 
     def name_to_idx(self, name: str) -> int:
+        "Returns the index of a link given its name."
         return self.link_names.index(name)
 
     def idx_to_name(self, idx: int, allow_world: bool = False) -> str:
+        """Returns the name of a link given its index. If `allow_world` is `True`,
+        returns `"world"` for index `-1`."""
         if allow_world and idx == -1:
             return "world"
         assert idx >= 0, "Worldbody index has no name."
         return self.link_names[idx]
 
     def idx_map(self, type: str) -> dict:
-        "type: is either `l` or `q` or `d`"
+        """Returns a dictionary mapping link names to their indices for a specified type
+        (`"l"`, `"q"`, or `"d"`)."""
         dict_int_slices = {}
 
         def f(_, idx_map, name: str, link_idx: int):
@@ -458,10 +724,11 @@ class System(_Base):
         return dict_int_slices
 
     def parent_name(self, name: str) -> str:
+        "Returns the name of the parent link for a given link."
         return self.idx_to_name(self.link_parents[self.name_to_idx(name)])
 
     def add_prefix(self, prefix: str = "") -> "System":
-        return self.replace(link_names=[prefix + name for name in self.link_names])
+        return self.add_prefix_suffix(prefix=prefix)
 
     def change_model_name(
         self,
@@ -469,6 +736,7 @@ class System(_Base):
         prefix: Optional[str] = None,
         suffix: Optional[str] = None,
     ) -> "System":
+        "Changes the name of the system model."
         if prefix is None:
             prefix = ""
         if suffix is None:
@@ -479,6 +747,7 @@ class System(_Base):
         return self.replace(model_name=name)
 
     def change_link_name(self, old_name: str, new_name: str) -> "System":
+        "Renames a specific link."
         old_idx = self.name_to_idx(old_name)
         new_link_names = self.link_names.copy()
         new_link_names[old_idx] = new_name
@@ -487,6 +756,7 @@ class System(_Base):
     def add_prefix_suffix(
         self, prefix: Optional[str] = None, suffix: Optional[str] = None
     ) -> "System":
+        "Adds either or, or both a prefix and suffix to all link names."
         if prefix is None:
             prefix = ""
         if suffix is None:
@@ -526,6 +796,7 @@ class System(_Base):
         return _update_sys_if_replace_joint_type(self, logic_replace_free_with_cor)
 
     def freeze(self, name: str | list[str]):
+        "Freezes the specified link(s), making them immovable (uses `frozen` joint)"
         if isinstance(name, list):
             sys = self
             for n in name:
@@ -544,6 +815,7 @@ class System(_Base):
         return _update_sys_if_replace_joint_type(self, logic_freeze)
 
     def unfreeze(self, name: str, new_joint_type: str):
+        "Unfreezes a frozen link and assigns it a new joint type."
         assert self.link_types[self.name_to_idx(name)] == "frozen"
         assert new_joint_type != "frozen"
 
@@ -560,7 +832,8 @@ class System(_Base):
         seed: int = 1,
         warn: bool = True,
     ):
-        "By default damping, stiffness are set to zero."
+        """Changes the joint type of a specified link.
+        By default damping, stiffness are set to zero."""
         from ring.algorithms import get_joint_model
 
         q_size, qd_size = Q_WIDTHS[new_joint_type], QD_WIDTHS[new_joint_type]
@@ -594,6 +867,7 @@ class System(_Base):
 
     @staticmethod
     def joint_type_simplification(typ: str) -> str:
+        "Returns a simplified name of the given joint type."
         if typ[:4] == "free":
             if typ == "free_2d":
                 return "free_2d"
@@ -608,23 +882,28 @@ class System(_Base):
 
     @staticmethod
     def joint_type_is_free_or_cor(typ: str) -> bool:
+        'Checks if a joint type is either "free" or "cor".'
         return System.joint_type_simplification(typ) in ["free", "cor"]
 
     @staticmethod
     def joint_type_is_spherical(typ: str) -> bool:
+        'Checks if a joint type is "spherical".'
         return System.joint_type_simplification(typ) == "spherical"
 
     @staticmethod
     def joint_type_is_free_or_cor_or_spherical(typ: str) -> bool:
+        'Checks if a joint type is "free", "cor", or "spherical".'
         return System.joint_type_is_free_or_cor(typ) or System.joint_type_is_spherical(
             typ
         )
 
     def findall_imus(self, names: bool = True) -> list[str] | list[int]:
+        "Finds all IMU sensors in the system."
         bodies = [name for name in self.link_names if name[:3] == "imu"]
         return bodies if names else [self.name_to_idx(n) for n in bodies]
 
     def findall_segments(self, names: bool = True) -> list[str] | list[int]:
+        "Finds all non-IMU segments in the system."
         imus = self.findall_imus(names=True)
         bodies = [name for name in self.link_names if name not in imus]
         return bodies if names else [self.name_to_idx(n) for n in bodies]
@@ -633,10 +912,12 @@ class System(_Base):
         return [self.idx_to_name(i) for i in bodies]
 
     def findall_bodies_to_world(self, names: bool = False) -> list[int] | list[str]:
+        "Returns all bodies directly connected to the world."
         bodies = [i for i, p in enumerate(self.link_parents) if p == -1]
         return self._bodies_indices_to_bodies_name(bodies) if names else bodies
 
     def find_body_to_world(self, name: bool = False) -> int | str:
+        "Returns the root body connected to the world."
         bodies = self.findall_bodies_to_world(names=name)
         assert len(bodies) == 1
         return bodies[0]
@@ -644,6 +925,7 @@ class System(_Base):
     def findall_bodies_with_jointtype(
         self, typ: str, names: bool = False
     ) -> list[int] | list[str]:
+        "Returns all bodies with the specified joint type."
         bodies = [i for i, _typ in enumerate(self.link_types) if _typ == typ]
         return self._bodies_indices_to_bodies_name(bodies) if names else bodies
 
@@ -781,20 +1063,25 @@ class System(_Base):
 
     @staticmethod
     def from_xml(path: str, seed: int = 1):
+        "Loads a system from an XML file."
         return ring.io.load_sys_from_xml(path, seed)
 
     @staticmethod
     def from_str(xml: str, seed: int = 1):
+        "Loads a system from an XML string."
         return ring.io.load_sys_from_str(xml, seed)
 
     def to_str(self, warn: bool = True) -> str:
+        "Serializes the system to an XML string."
         return ring.io.save_sys_to_str(self, warn=warn)
 
     def to_xml(self, path: str) -> None:
+        "Saves the system to an XML file."
         ring.io.save_sys_to_xml(self, path)
 
     @classmethod
     def create(cls, path_or_str: str, seed: int = 1) -> "System":
+        "Creates a `System` instance from an XML file or string."
         path = Path(path_or_str).with_suffix(".xml")
 
         exists = False
@@ -814,7 +1101,8 @@ class System(_Base):
         q: jax.Array,
         custom_joints: dict[str, Callable] = {},
     ) -> jax.Array:
-        """Map a coordinate vector `q` to the minimal coordinates vector of the sys"""
+        """Converts a coordinate vector to minimal coordinates (`q`), applying
+        constraints such as quaternion normalization."""
         # Does, e.g.
         # - normalize quaternions
         # - hinge joints in [-pi, pi]
@@ -1026,14 +1314,37 @@ def _scan_sys(sys: System, f: Callable, in_types: str, *args, reverse: bool = Fa
 
 @struct.dataclass
 class State(_Base):
-    """The static and dynamic state of a system in minimal and maximal coordinates.
-    Use `.create()` to create this object.
-
-    Args:
-        q (jax.Array): System state in minimal coordinates (equals `sys.q_size()`)
-        qd (jax.Array): System velocity in minimal coordinates (equals `sys.qd_size()`)
-        x: (Transform): Maximal coordinates of all links. From epsilon-to-link.
     """
+    Represents the state of a dynamic system in minimal and maximal coordinates.
+
+    The `State` class encapsulates both the configuration (`q`) and velocity (`qd`)
+    of the system in minimal coordinates, as well as the corresponding transforms (`x`)
+    in maximal coordinates.
+
+    Attributes:
+        q (jax.Array):
+            The joint positions (generalized coordinates) of the system. The size
+            of `q` matches `sys.q_size()`.
+        qd (jax.Array):
+            The joint velocities (generalized velocities) of the system. The size
+            of `qd` matches `sys.qd_size()`.
+        x (Transform):
+            The maximal coordinate representation of all system links, expressed as
+            a `Transform` object.
+
+    Methods:
+        create(sys: System, q: Optional[jax.Array] = None,
+               qd: Optional[jax.Array] = None, x: Optional[Transform] = None,
+               key: Optional[jax.Array] = None,
+               custom_joints: dict[str, Callable] = {}) -> State:
+            Creates a `State` instance for a given system with optional initial conditions.
+
+    Usage:
+        >>> sys = System.create("model.xml")
+        >>> state = State.create(sys)
+        >>> print(state.q.shape)  # Should match sys.q_size()
+        >>> print(state.qd.shape)  # Should match sys.qd_size()
+    """  # noqa: E501
 
     q: jax.Array
     qd: jax.Array
@@ -1048,19 +1359,37 @@ class State(_Base):
         x: Optional[Transform] = None,
         key: Optional[jax.Array] = None,
         custom_joints: dict[str, Callable] = {},
-    ):
-        """Create state of system.
+    ) -> "State":
+        """
+        Creates a `State` instance for the given system with optional initial conditions.
+
+        If no initial values are provided, joint positions (`q`) and velocities (`qd`)
+        are initialized to zero, except for free and spherical joints, which have unit quaternions.
 
         Args:
-            sys (System): The system for which to create a state.
-            q (jax.Array, optional): The joint values of the system. Defaults to None.
-            Which then defaults to zeros.
-            qd (jax.Array, optional): The joint velocities of the system.
-            Defaults to None. Which then defaults to zeros.
+            sys (System):
+                The system for which to create a state.
+            q (Optional[jax.Array], default=None):
+                Initial joint positions. If `None`, defaults to zeros, with unit quaternion initialization
+                for free and spherical joints.
+            qd (Optional[jax.Array], default=None):
+                Initial joint velocities. If `None`, defaults to zeros.
+            x (Optional[Transform], default=None):
+                Initial maximal coordinates of the system links. If `None`, defaults to zero transforms.
+            key (Optional[jax.Array], default=None):
+                Random key for initializing `q` if no values are provided.
+            custom_joints (dict[str, Callable], default={}):
+                Custom joint functions for mapping coordinate vectors to minimal coordinates.
 
         Returns:
-            (State): Create State object.
-        """
+            State: A new instance of the `State` class representing the initialized system state.
+
+        Example:
+            >>> sys = System.create("model.xml")
+            >>> state = State.create(sys)
+            >>> print(state.q.shape)  # Should match sys.q_size()
+            >>> print(state.qd.shape)  # Should match sys.qd_size()
+        """  # noqa: E501
         if key is not None:
             assert q is None
             q = jax.random.normal(key, shape=(sys.q_size(),))

ring/ml/base.py

@@ -297,7 +297,8 @@ class NoGraph_FilterWrapper(AbstractFilterWrapper):
 
         if self._quat_normalize:
             assert yhat.shape[-1] == 4, f"yhat.shape={yhat.shape}"
-            yhat = ring.maths.safe_normalize(yhat)
+            # yhat = ring.maths.safe_normalize(yhat)
+            yhat = yhat / jnp.linalg.norm(yhat, axis=-1, keepdims=True)
 
         return yhat, state