Function serialize_qp

diffmimic/utils/io.py:17–28

(qp)

def serialize_qp(qp) -> jnp.array:
    pos = qp.pos
    rot = qp.rot
    vel = qp.vel
    ang = qp.ang
    batch_dim = pos.shape[:-2]
    nparray = []
    nparray.append(pos.reshape(batch_dim + (-1,)))
    nparray.append(rot.reshape(batch_dim + (-1,)))
    nparray.append(vel.reshape(batch_dim + (-1,)))
    nparray.append(ang.reshape(batch_dim + (-1,)))
    return jnp.concatenate(nparray, axis=-1)