jax-shac

• An implementation of Short Horizon Actor Critic (Xu; 2022) writen in Jax
• Simulation using the Mujoco MJX simulator

Results

Inverted Pendulum

Run Time: 1 min jit, 2 min training

Known Issues: For some random seeds, you get drift in the cart position.

1 DOF Hopper

Run Time: 1 min jit, 2 min training

Known Issues: As seen in the rewards figure, training can be unstable.

Warning: MJX + Exploding Gradients

• Having great difficulty applying SHAC to get Anymal to walk with default 32-bit precision. (See Mujoco for an example with 64-bit precision)
• Hypothesis: it's because quadruped gait is very contact-rich, leading to uninformative gradients.

32-step rollout. Ground flashes red when step jacobian is greater than 10e2.

Setup

• pip install -r requirements.txt
• Add the parent folder of this repository to your PYTHONPATH environment variable.