Well done! At this point you installed coax and JAX, you wrote a linear regression model in JAX and Haiku, and trained your own RL agents.
The best place to continue is to head over to the Examples page to have a look at other agents. For instance, you may be interested in training an agent for continuous action spaces. In which case, you’ll like this example: Pendulum with DDPG.
The world’s your oyster!