Simpler indexing in JAX DQN

[joelberkeley]

Description

Since this is a "clean" RL repo I thought it might help to improve readability. If this PR is too small to be worthwhile, that's fine.

I don't believe this should affect performance or functionality.

I've not read the contributing guidelines properly. If that's a blocker, feel free to reject PR.

I didn't work out how to get the dependencies to work, so I didn't re-run the file as a smoke test

Types of changes

Readability - none of the below

• Bug fix
• New feature
• New algorithm
• Documentation

Checklist:

• I've read the CONTRIBUTION guide (required).
• I have ensured pre-commit run --all-files passes (required).
• I have updated the tests accordingly (if applicable).
• I have updated the documentation and previewed the changes via mkdocs serve.
  • I have explained note-worthy implementation details.
  • I have explained the logged metrics.
  • I have added links to the original paper and related papers.

[vercel]

@joelberkeley is attempting to deploy a commit to the Costa Huang's projects Team on Vercel.

A member of the Team first needs to authorize it.