reference_generator refactoring

[LockedFlysher]

Problem Description

in mpc_utils.py

hip = jnp.tile(p, n_contact) + foot0 @ jax.scipy.linalg.block_diag(*([Ryaw] * n_contact)).T

The variable is named hip but actually computes foot positions in global coordinates, not hip joint positions. This causes:

• Semantic mismatch and confusion
• Poor code readability
• Maintenance difficulties

What the calculation does

• p: Robot base position in global frame
• foot0: Foot positions in local coordinates
• Ryaw: Base yaw rotation matrix
• Result: Foot positions in global frame

Proposed Solution

Rename the variable to accurately reflect its purpose:

# Recommended
foot_global_pos = jnp.tile(p, n_contact) + foot0 @ jax.scipy.linalg.block_diag(*([Ryaw] * n_contact)).T

# Alternative
feet_world_positions = jnp.tile(p, n_contact) + foot0 @ jax.scipy.linalg.block_diag(*([Ryaw] * n_contact)).T

Labels

bug code quality good first issue

[lorenzo96-cmd]

Thanks for your interest in the repo. Sorry for the confusion, hip was definitely a misleading name. I've changed it to foot0_projected, as the operation you pointed out projects the default foot position from the base frame to a global horizontal frame.

Best,
Lorenzo

[LockedFlysher]

It's hard for engnieers and scientists to keep the code clean and easy to understand, I'm very happy to meet the work like yours. I can share my edited version, haha

def reference_generator(
        #fixed params
        use_terrain_estimator, N, dt, n_joints, n_contact, mass, foot0, q0,
        #runtime variables
        t_timer, x, foot, input, duty_factor, step_freq, step_height, liftoff, contact,clearence_speed):
    """Generates reference trajectories for quadruped robot locomotion.

    This function generates comprehensive reference trajectories including base position/orientation,
    joint angles, foot positions, contact sequences, and ground reaction forces for a quadruped
    robot over a prediction horizon.

    Args:
        use_terrain_estimator (bool): Whether to use terrain estimator for attitude reference.
        N (int): Prediction horizon length.
        dt (float): Time step size in seconds.
        n_joints (int): Number of robot joints.
        n_contact (int): Number of contact points (typically 4 for quadruped).
        mass (float): Total robot mass in kg.
        foot0 (jnp.ndarray): Initial foot positions in body frame, shape (3*n_contact,).
        q0 (jnp.ndarray): Initial joint angles, shape (n_joints,).
        t_timer (jnp.ndarray): Current gait timers for each foot, shape (n_contact,).
        x (jnp.ndarray): Current robot state vector.
        foot (jnp.ndarray): Current foot positions in world frame, shape (3*n_contact,).
        input (jnp.ndarray): Command inputs [vx, vy, vz, wx, wy, wz, height].
        duty_factor (float): Gait duty factor (stance time / cycle time).
        step_freq (float): Step frequency in Hz.
        step_height (float): Maximum foot clearance height during swing.
        liftoff (jnp.ndarray): Last liftoff positions for each foot, shape (3*n_contact,).
        contact (jnp.ndarray): Current contact states, shape (n_contact,).
        clearence_speed (float): Initial velocity for foot clearance trajectory.

    Returns:
        tuple: A tuple containing:
            - full_reference (jnp.ndarray): Complete reference trajectory, shape (N+1, state_dim).
            - contact_output (jnp.ndarray): Contact sequence, shape (N+1, n_contact).
            - liftoff (jnp.ndarray): Updated liftoff positions, shape (3*n_contact,).
    """

    # ============ 1. Extract Current State Information ============
    p = x[:3]  # Position [x, y, z]
    quat = x[3:7]  # Quaternion [w, x, y, z]
    dp = x[7 + n_joints:10 + n_joints]  # Linear velocity [vx, vy, vz]

    # ========== Complete Rotation Matrix Calculation Function ==========
    def quaternion_to_rotation_matrix(q):
        """Converts quaternion to rotation matrix.

        Args:
            q (jnp.ndarray): Quaternion [w, x, y, z].

        Returns:
            jnp.ndarray: 3x3 rotation matrix.
        """
        w, x, y, z = q[0], q[1], q[2], q[3]
        return jnp.array([
            [1 - 2 * (y ** 2 + z ** 2), 2 * (x * y - w * z), 2 * (x * z + w * y)],
            [2 * (x * y + w * z), 1 - 2 * (x ** 2 + z ** 2), 2 * (y * z - w * x)],
            [2 * (x * z - w * y), 2 * (y * z + w * x), 1 - 2 * (x ** 2 + y ** 2)]
        ])

    # Calculate complete rotation matrix (including roll, pitch, yaw)
    R_body = quaternion_to_rotation_matrix(quat)

    # Calculate yaw angle and rotation matrix for terrain estimator
    # TODO: Check if these parameters are needed, might be unused but kept for terrain estimator
    yaw = jnp.arctan2(2 * (quat[0] * quat[3] + quat[1] * quat[2]),
                      1 - 2 * (quat[2] * quat[2] + quat[3] * quat[3]))
    Ryaw = jnp.array([[jnp.cos(yaw), -jnp.sin(yaw), 0],
                      [jnp.sin(yaw), jnp.cos(yaw), 0],
                      [0, 0, 1]])

    # Calculate body height: commanded height + average height of feet at liftoff moments
    proprio_height = input[6] + jnp.sum(liftoff[2::3]) / n_contact
    p = jnp.array([p[0], p[1], proprio_height])

    # ============ 2. Generate Base Attitude Reference Trajectory ============
    if use_terrain_estimator:
        # Generate attitude based on terrain estimation
        quat_ref = jnp.tile(terrain_orientation(liftoff, Ryaw), (N + 1, 1))
    else:
        # Maintain horizontal attitude
        quat_ref = jnp.tile(jnp.array([1, 0, 0, 0]), (N + 1, 1))

    # Joint angles maintain initial values
    q_ref = jnp.tile(q0, (N + 1, 1))

    # ============ 3. Generate Position and Velocity Reference Trajectories ============
    # Transform body-frame velocity to world frame
    ref_lin_vel = R_body @ input[:3]
    ref_ang_vel = input[3:6]

    # Generate position reference trajectory (integrate along velocity direction)
    p_ref_x = jnp.arange(N + 1) * dt * ref_lin_vel[0] + p[0]
    p_ref_y = jnp.arange(N + 1) * dt * ref_lin_vel[1] + p[1]
    p_ref_z = jnp.ones(N + 1) * proprio_height
    p_ref = jnp.stack([p_ref_x, p_ref_y, p_ref_z], axis=1)

    # Velocity reference trajectories remain constant
    dp_ref = jnp.tile(ref_lin_vel, (N + 1, 1))
    omega_ref = jnp.tile(ref_ang_vel, (N + 1, 1))

    # ============ 4. Initialize Foot-Related Variables ============
    # Initialize foot position references
    foot_ref = jnp.tile(foot, (N + 1, 1))

    # ========== Correction: Calculate foot global positions using complete rotation matrix ==========
    # Calculate foot positions in world frame (considering roll, pitch, yaw)
    foot_global_pos = (jnp.tile(p, n_contact) +
                       foot0 @ jax.scipy.linalg.block_diag(*([R_body] * n_contact)).T)

    # ============ 5. Early Contact Detection ============
    des_contact, current_timer = timer_run(duty_factor, step_freq, t_timer, dt)
    # Detect unexpected contact in late swing phase (adjustment for higher terrain)
    early_contact = jnp.where(
        jnp.logical_and(
            jnp.logical_and(des_contact == 0, contact == 1),
            current_timer > 0.5 + 0.5 * duty_factor
        ), 1, 0)

    # ============ 6. Main Loop for Foot Trajectory Generation ============
    def foot_fn(t, carry):
        """Main function for foot trajectory generation at each time step.

        Args:
            t (int): Current time step index.
            carry (tuple): Carry variables from previous iteration.

        Returns:
            tuple: Updated carry variables for next iteration.
        """
        # Extract carry variables
        timer_seq, contact_sequence, new_foot, liftoff_x, liftoff_y, liftoff_z, grf_new = carry

        # Extract previous foot positions
        prev_foot_x = new_foot[t - 1, ::3]
        prev_foot_y = new_foot[t - 1, 1::3]
        prev_foot_z = new_foot[t - 1, 2::3]

        # Update timer and contact sequence
        new_contact_sequence, new_timer = timer_run(duty_factor, step_freq, timer_seq[t - 1, :], dt)
        contact_sequence = contact_sequence.at[t, :].set(new_contact_sequence)
        timer_seq = timer_seq.at[t, :].set(new_timer)

        # -------- Detect Liftoff Moment and Update Liftoff Positions --------
        # Condition: currently not in contact AND was in contact at previous step = liftoff moment
        liftoff_condition = jnp.logical_and(
            jnp.logical_not(contact_sequence[t, :]),  # Currently not in contact
            contact_sequence[t - 1, :]  # Was in contact at previous step
        )

        # Record foot position at liftoff moment
        liftoff_x = jnp.where(liftoff_condition, prev_foot_x, liftoff_x)
        liftoff_y = jnp.where(liftoff_condition, prev_foot_y, liftoff_y)
        liftoff_z = jnp.where(liftoff_condition, prev_foot_z, liftoff_z)

        # -------- Calculate Foothold Positions --------
        def calc_foothold(direction):
            """Calculate desired foothold position for a given direction.

            Args:
                direction (int): Direction index (0=x, 1=y, 2=z).

            Returns:
                jnp.ndarray: Desired foothold positions for all feet.
            """
            # Feedforward term: prediction based on desired velocity
            feedforward = 0.5 * ref_lin_vel[direction] * duty_factor / step_freq
            # Feedback term: correction based on velocity error
            feedback = jnp.sqrt(input[6] / 9.81) * (dp[direction] - ref_lin_vel[direction])
            # Final foothold = current foot position + feedforward + feedback
            return foot_global_pos[direction::3] + feedforward + feedback

        # -------- Spline Trajectory Function Definitions --------
        def cubic_splineXY(start_pos, end_pos, initial_velocity, t_normalized):
            """Cubic spline trajectory for XY directions.

            Args:
                start_pos (jnp.ndarray): Starting positions.
                end_pos (jnp.ndarray): Ending positions.
                initial_velocity (float): Initial velocity.
                t_normalized (jnp.ndarray): Normalized time parameter [0, 1].

            Returns:
                jnp.ndarray: Interpolated positions.
            """
            a0 = start_pos
            a1 = initial_velocity
            a2 = 3 * (end_pos - start_pos) - 2 * initial_velocity
            a3 = initial_velocity - 2 * (end_pos - start_pos)
            return a0 + a1 * t_normalized + a2 * t_normalized ** 2 + a3 * t_normalized ** 3

        def cubic_splineZ(start_pos, end_pos, max_height, t_normalized):
            """Quartic spline trajectory for Z direction (parabolic shape).

            Args:
                start_pos (jnp.ndarray): Starting positions.
                end_pos (jnp.ndarray): Ending positions.
                max_height (jnp.ndarray): Maximum clearance height.
                t_normalized (jnp.ndarray): Normalized time parameter [0, 1].

            Returns:
                jnp.ndarray: Interpolated positions.
            """
            initial_speed = 0.7
            a = 16 * max_height - 8 * end_pos - 8 * start_pos - 2 * initial_speed
            b = 5 * initial_speed + 14 * end_pos + 18 * start_pos - 32 * max_height
            c = 16 * max_height - 5 * end_pos - 11 * start_pos - 4 * initial_speed
            d = initial_speed
            e = start_pos
            return a * t_normalized ** 4 + b * t_normalized ** 3 + c * t_normalized ** 2 + d * t_normalized + e

        # -------- Calculate Swing Trajectory --------
        # Initial swing velocity (opposite to motion direction)
        initial_speed = -ref_lin_vel / (jnp.linalg.norm(ref_lin_vel) + 1e-6) * clearence_speed
        # Normalized time parameter for swing phase
        swing_parameter = (new_timer - duty_factor) / (1 - duty_factor)
        # Calculate foothold positions for all feet
        foothold_x = calc_foothold(0)
        foothold_y = calc_foothold(1)

        # Select foot positions based on contact state
        # Contact state or early contact: maintain previous position
        # Swing state: use spline trajectory
        new_foot_x = jnp.where(
            jnp.logical_or(new_contact_sequence > 0, early_contact == 1),
            prev_foot_x,  # Maintain position during contact
            cubic_splineXY(liftoff_x, foothold_x, initial_speed[0], swing_parameter)  # Swing trajectory
        )

        new_foot_y = jnp.where(
            jnp.logical_or(new_contact_sequence > 0, early_contact == 1),
            prev_foot_y,
            cubic_splineXY(liftoff_y, foothold_y, initial_speed[1], swing_parameter)
        )

        new_foot_z = jnp.where(
            jnp.logical_or(new_contact_sequence > 0, early_contact == 1),
            prev_foot_z,
            cubic_splineZ(liftoff_z, liftoff_z, liftoff_z + step_height, swing_parameter)
        )

        # Update foot positions
        new_foot = new_foot.at[t, ::3].set(new_foot_x)
        new_foot = new_foot.at[t, 1::3].set(new_foot_y)
        new_foot = new_foot.at[t, 2::3].set(new_foot_z)

        # -------- Calculate Ground Reaction Forces --------
        # Simple allocation: distribute total weight evenly among contacting feet
        grf_z = new_contact_sequence * mass * 9.81 / (jnp.sum(new_contact_sequence) + 1e-5)
        grf_new = grf_new.at[t, 2::3].set(grf_z)

        return (timer_seq, contact_sequence, new_foot, liftoff_x, liftoff_y, liftoff_z, grf_new)

    # ============ 7. Execute Main Loop ============
    # Separate liftoff coordinates
    liftoff_x = liftoff[::3]
    liftoff_y = liftoff[1::3]
    liftoff_z = liftoff[2::3]

    # Initialize loop variables
    contact_sequence = jnp.zeros((N + 1, n_contact))  # Contact sequence (all initially non-contact)
    grf_ref = jnp.zeros((N + 1, 3 * n_contact))  # Ground reaction forces
    timer_sequence_in = jnp.tile(t_timer, (N + 1, 1))  # Timer sequence with shape (N+1) x n_contact

    init_carry = (timer_sequence_in, contact_sequence, foot_ref, liftoff_x, liftoff_y, liftoff_z, grf_ref)

    # Execute loop to generate trajectories using JAX's fori_loop for fixed iterations
    final_carry = jax.lax.fori_loop(0, N + 1, foot_fn, init_carry)
    timer_sequence, contact_sequence, foot_ref, liftoff_x, liftoff_y, liftoff_z, grf_ref = final_carry

    # ============ 8. Update Output Liftoff ============
    liftoff = liftoff.at[::3].set(liftoff_x)
    liftoff = liftoff.at[1::3].set(liftoff_y)
    liftoff = liftoff.at[2::3].set(liftoff_z)

    # ============ 9. Assemble Output ============
    # Complete reference trajectory, each row represents full state vector at one time step
    full_reference = jnp.concatenate([
        p_ref,  # Position reference
        quat_ref,  # Attitude reference
        q_ref,  # Joint angle reference
        dp_ref,  # Linear velocity reference
        omega_ref,  # Angular velocity reference
        foot_ref,  # Foot position reference
        contact_sequence,  # Contact sequence
        grf_ref  # Ground reaction force reference
    ], axis=1)

    # Contact sequence as separate output
    contact_output = jnp.concatenate([contact_sequence], axis=1)

    return full_reference, contact_output, liftoff

[LockedFlysher]

By the way, I have made a change in this code, and it's so good to share with you:
The full rotation matrix: I found the 'ref_lin_vel' and 'foot0_projected' are not using the full rotation matrix, which may cause trouble.

[lorenzo96-cmd]

ref_lin_vel and foot0_projected are computed correctly, as the input reference velocity is given in the local horizontal frame. This frame is commonly used in robotics: the yaw is aligned with the robot’s orientation, while the pitch and roll remain aligned with the world frame. For this reason, when computing ref_lin_vel from the input, we apply only the yaw rotation (R_yaw). As mentioned earlier, foot0_projected is also a projection from the horizontal frame.