| 59 | world_ang = ase_ang[i][1:] |
| 60 | |
| 61 | def init(qp): |
| 62 | pos = jp.index_update(qp.pos, 0, pelvis_trans[i]) |
| 63 | rot = ase_poses[i][0] / jp.norm(ase_poses[i][0]) # important |
| 64 | # print(rot) |
| 65 | rot = math.quat_mul(math.euler_to_quat(np.array([0., -90, 0.])), rot) |
| 66 | # print(rot) |
| 67 | # raise |
| 68 | rot = jp.index_update(qp.rot, 0, rot) |
| 69 | vel = jp.index_update(qp.vel, 0, ase_vel[i][0]) |
| 70 | ang = jp.index_update(qp.ang, 0, ase_ang[i][0]) |
| 71 | qp = qp.replace(pos=pos, rot=rot, vel=vel, ang=ang) |
| 72 | return qp |
| 73 | |
| 74 | qp = init(qp) |
| 75 | amp_rot = local_rot[joint_order] |