NVlabs / trajdata

A unified interface to many trajectory forecasting datasets.

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

nuPlan already standardizes dynamic quantities

BorisIvanovic opened this issue · comments

As can be seen in their schema, velocities and accelerations are already standardized for the ego-vehicle in nuPlan, meaning when trajdata tries to standardize things here, it yields incorrect velocities/accelerations.

This is unique to nuPlan as they handle the relativization of states for users (in contrast to other datasets). Will have to think about the best way to handle this prior to fixing it.

As a temporary fix, how about transforming the velocities and accelerations back to the world in the nuplan/nuplan_dataset.py?

def transform_v_a_to_world(ego_df):
    # Extract yaw angles (headings) and convert to NumPy array
    headings = np.array(ego_df['heading'])
    
    # Create a batch_dims variable; it's assumed that your ego_df is 1D along the 'heading' axis
    batch_dims = [len(ego_df)]
    
    # Create rotation matrix for each heading
    cos_headings = np.cos(headings)
    sin_headings = np.sin(headings)
    
    rotation_matrix = np.array([
        [cos_headings, -sin_headings],
        [sin_headings, cos_headings]
    ]).transpose(2, 0, 1)

    # Assuming ego_data contains ['vx', 'vy', 'ax', 'ay']
    ego_data = ego_df[['vx', 'vy', 'ax', 'ay']].values
    batch_dims = ego_data.shape[0]

    # Reshape ego_data for multiplication, aligning it to (batch_size, 2, 2)
    ego_data_reshaped = ego_data.reshape(batch_dims, 2, 2)

    # Perform matrix multiplication
    transformed_data = (ego_data_reshaped @ rotation_matrix).reshape(batch_dims, 4)

    # Assign back to DataFrame
    ego_df[['vx', 'vy', 'ax', 'ay']] = transformed_data

    return ego_df