🪶 Zig-Zag trajectory generation

[jc-bao]

Generate a Python script with docstring to implement the following function. The function is used to generate a 2D reference trajectory with velocity.

Input: timestep dt, max trajectory length traj_len and desired velocity v_tar. Output: trajectory y coordinate traj_y, z coordinate traj_z and velocity traj_y_dot, traj_z_dot

1. randomly pick up 5 points in 2D y-z space within [-1.5, 1.5], called p[0],..., p[4]
2. generate a periodic trajectory by sampling point from path p[0]->p[1]->...->p[4]->p[0]->... For every two points in the trajectory, their distance is v_tar*dt.
3. calculate the velocity of each point by differentiating the trajectory.

[jc-bao]

🤖

import numpy as np

def generate_2d_reference_trajectory(dt, traj_len, v_tar):
    """
    Generate a 2D reference trajectory with velocity.

    Parameters:
    dt (float): Timestep.
    traj_len (int): Maximum trajectory length.
    v_tar (float): Desired velocity.

    Returns:
    tuple: A tuple containing arrays for trajectory y coordinate (traj_y), 
           z coordinate (traj_z) and velocity (traj_y_dot, traj_z_dot)
    """

    # 1. Randomly pick up 5 points in 2D y-z space within [-1.5, 1.5]
    points = np.random.uniform(-1.5, 1.5, (5, 2))

    # 2. Generate a periodic trajectory by sampling points from the path
    traj_y = []
    traj_z = []
    distance = v_tar * dt
    num_segments = points.shape[0]

    for i in range(traj_len):
        start_point = points[i % num_segments]
        end_point = points[(i + 1) % num_segments]

        # Calculate direction vector and normalize
        direction = end_point - start_point
        direction /= np.linalg.norm(direction)

        # Calculate new point by moving in the direction with the desired distance
        new_point = start_point + direction * distance

        traj_y.append(new_point[0])
        traj_z.append(new_point[1])

        # Update start_point for the next iteration
        points[i % num_segments] = new_point

    traj_y = np.array(traj_y)
    traj_z = np.array(traj_z)

    # 3. Calculate the velocity of each point by differentiating the trajectory
    traj_y_dot = np.gradient(traj_y, dt)
    traj_z_dot = np.gradient(traj_z, dt)

    return traj_y, traj_z, traj_y_dot, traj_z_dot

[jc-bao]

🗣️

You have made a mistake. For each trajectory segment, you start from the same start_point, which means the trajectory stays at the same point over time. The objective of the function is to generate a zig-zag trajectory. Let's consider the problem in this way:

1. you can interpolate between two points to get a linear trajectory.
2. repeat this process to get all linear trajectory points and connect them together.

[jc-bao]

🤖

import numpy as np

def generate_2d_reference_trajectory(dt, traj_len, v_tar):
    """
    Generate a 2D reference trajectory with velocity.

    Parameters:
    dt (float): Timestep.
    traj_len (int): Maximum trajectory length.
    v_tar (float): Desired velocity.

    Returns:
    tuple: A tuple containing arrays for trajectory y coordinate (traj_y), 
           z coordinate (traj_z) and velocity (traj_y_dot, traj_z_dot)
    """

    # 1. Randomly pick up 5 points in 2D y-z space within [-1.5, 1.5]
    points = np.random.uniform(-1.5, 1.5, (5, 2))

    # 2. Generate a zig-zag trajectory by interpolating between points
    traj_y = []
    traj_z = []
    distance = v_tar * dt
    num_segments = points.shape[0]

    current_point = points[0]
    next_point_idx = 1

    for _ in range(traj_len):
        next_point = points[next_point_idx]

        # Calculate direction vector and normalize
        direction = next_point - current_point
        direction /= np.linalg.norm(direction)

        # Calculate new point by moving in the direction with the desired distance
        new_point = current_point + direction * distance

        traj_y.append(new_point[0])
        traj_z.append(new_point[1])

        current_point = new_point

        # Check if we have reached the next point
        if np.linalg.norm(next_point - current_point) < distance:
            next_point_idx = (next_point_idx + 1) % num_segments

    traj_y = np.array(traj_y)
    traj_z = np.array(traj_z)

    # 3. Calculate the velocity of each point by differentiating the trajectory
    traj_y_dot = np.gradient(traj_y, dt)
    traj_z_dot = np.gradient(traj_z, dt)

    return traj_y, traj_z, traj_y_dot, traj_z_dot

[jc-bao]

🗣️

The GPT has got things wrong. Let's fix that.

Generate a Python script based on Jax with docstring.

Input: point1 in y-z plane, point2 in y-z plane. Point number. Output: linear interpolation between two points in the y and z axis.

Avoid using for loop for better performance.

[jc-bao]

🤖

    # Generate an array of evenly spaced values between 0 and 1 with shape (num_points, 1)
    t = jnp.linspace(0, 1, num_points).reshape(-1, 1)

    # Perform linear interpolation along the y and z axes
    interpolated_y = y1 + (y2 - y1) * t
    interpolated_z = z1 + (z2 - z1) * t

[jc-bao]

🗣️

Generate six random points in 2D y-z space with Python and jax.

1. sample the first point uniformly from [-1.5, 1.5] (y[0], z[0]).
2. sample the next point with a distance sampling from [1.5, 2.5]. To make sure the point is not too far away from the centre, the displacement angle should be around the previous point angle to the centre jnp.atans(y[0],z[0]). Add a little displacement to that angle which sample from [-jnp.pi/4, jnp.pi/4] too add a little randomness.
3. repeat step 2 to get all 6 points.

[jc-bao]

🤖

import jax.numpy as jnp
import jax
from jax.random import PRNGKey, uniform
import matplotlib.pyplot as plt

def generate_points(num_points=6):
    points = jnp.zeros((num_points, 2))
    key = PRNGKey(3)

    # Sample the first point uniformly from [-1.5, 1.5]
    key, subkey = jax.random.split(key)
    y = uniform(subkey, minval=-1.5, maxval=1.5)
    key, subkey = jax.random.split(key)
    z = uniform(subkey, minval=-1.5, maxval=1.5)
    points = points.at[0].set(jnp.array([y, z]))

    for i in range(1, num_points):
        # Calculate the previous point angle to the center
        prev_angle = jnp.arctan2(points[i-1, 1], points[i-1, 0]) + jnp.pi

        # Sample a random displacement angle from [-pi/4, pi/4]
        key, subkey = jax.random.split(key)
        delta_angle = uniform(subkey, minval=-jnp.pi/4, maxval=jnp.pi/4)

        # Calculate the new angle
        angle = prev_angle + delta_angle

        # Sample the distance from [1.5, 2.5]
        key, subkey = jax.random.split(key)
        distance = uniform(subkey, minval=1.5, maxval=2.5)

        # Calculate the new point
        y = points[i-1, 0] + distance * jnp.cos(angle)
        z = points[i-1, 1] + distance * jnp.sin(angle)
        points = points.at[i].set(jnp.array([y, z]))

    return points

points = generate_points()

# Create the plot
plt.figure()
for i in range(6):
    plt.scatter(points[i, 0], points[i, 1], marker='o', label=f'{i}')
    # plot arrow from previous point to current point
    if i > 0:
        plt.arrow(points[i-1, 0], points[i-1, 1], points[i, 0] - points[i-1, 0], points[i, 1] - points[i-1, 1], 
                  color='black', width=0.01, head_width=0.05, length_includes_head=True, alpha=i / 6)
plt.legend()
plt.scatter(0, 0, marker='x', color='red', label='Center')

# Add labels and legend
plt.xlabel('y')
plt.ylabel('z')
plt.title('Random Points in 2D y-z Space')
plt.legend()

# Set equal aspect ratio and display the plot
plt.gca().set_aspect('equal', adjustable='box')
plt.show()