def rotation_matrix_to_ortho6d_1d(matrix: tf.Tensor) -> tf.Tensor:
    """
    The orhto6d represents the first two column vectors a1 and a2 of the
    rotation matrix: [ | , |,  | ]
                     [ a1, a2, a3]
                     [ | , |,  | ]
    Input: (3, 3)
    Output: (6,)
    This function is used to make tensorflow happy.
    """
    ortho6d = matrix[:, :2]
    # Transpose the last two dimension
    ortho6d = tf.transpose(ortho6d)
    # Flatten the last two dimension
    ortho6d = tf.reshape(ortho6d, [6])
    return ortho6d

    The orhto6d represents the first two column vectors a1 and a2 of the
    rotation matrix: [ | , |,  | ]
                     [ a1, a2, a3]
                     [ | , |,  | ]
    Input: (3, 3)
    Output: (6,)
    This function is used to make tensorflow happy.
    
policy/RDT/data/utils.py