actor_log_prob.transpose in MPO causes error

[zhaoyi11]

Hi,

Thanks so much for the clean implementation. In MPO, using actor_log_prob.transpose((0, 1)) causes error.

jax-rl/jax_rl/MPO.py

Line 236 in a2c6067

actor_log_prob = actor_log_prob.transpose((0, 1))

TypeError: transpose permutation isn't a permutation of operand dimensions, got permutation (0, 1) for operand shape (256, 20, 1).

I checked the shape of actor_log_prob and weights, they are in the same shape. After removing this line, I can get the samilar performance in cartpole_swingup env as shown in README.md.

Best,
Yi

[henry-prior]

Looks like this line was unnecessary, removed it in #11