def rmat_to_rot6(rmat):
  """Projects rotation matrix to 6-dim "Gram-Schmidt-able" representation.

  The "rot6" representation is a 6-DOF representation of an orientation that is
  homeomorphic with SO(3). It is not minimal like an euler or axis-angle, but
  it is smooth over the full range of rotations, unlike eulers, quaternions, and
  axis-angle representations. See the original paper for details:
    "On the Continuity of Rotation Representations in Neural Networks"
    https://arxiv.org/pdf/1812.07035.pdf

  Args:
    rmat: A 3x3 rotation matrix.

  Returns:
    A 6-dimensional array containing the first two columns of `rmat`.  This
    representation can be mapped back to `rmat` using `rot6_to_rmat`.
  """
  return rmat[:, :-1].T.flatten()