Jax DQN & Improvements to AgentStats

[omardrwch]

• Implementation of a simple DQN agent in JAX;
• Improvements to AgentStats class, so that it works well with JAX agents:
  • Using spawn instead of fork in multiprocessing (see jax does not work with the multiprocessing "fork" strategy. jax-ml/jax#1805);
  • Fixed locks used by threads and processes.
• An evaluation environment eval_env is now created in the Agent class (and can be optionally sent to the constructor). This makes it easier to define the Agent.eval() method.

[codecov]

Codecov Report

Merging #51 (806a726) into main (1e36b18) will decrease coverage by 0.15%.
The diff coverage is 80.68%.

@@            Coverage Diff             @@
##             main      #51      +/-   ##
==========================================
- Coverage   87.05%   86.89%   -0.16%     
==========================================
  Files         142      149       +7     
  Lines        7186     7457     +271     
==========================================
+ Hits         6256     6480     +224     
- Misses        930      977      +47     

Impacted Files	Coverage Δ
setup.py	0.00% <ø> (ø)
rlberry/agents/jax/dqn/dqn.py	73.05% <73.05%> (ø)
rlberry/stats/agent_stats.py	83.65% <73.68%> (+0.96%)	⬆️
rlberry/envs/utils.py	89.47% <89.47%> (ø)
rlberry/agents/jax/utils/replay_buffer.py	90.62% <90.62%> (ø)
rlberry/agents/agent.py	77.77% <100.00%> (+0.76%)	⬆️
rlberry/agents/jax/__init__.py	100.00% <100.00%> (ø)
rlberry/agents/jax/nets/common.py	100.00% <100.00%> (ø)
rlberry/agents/jax/tests/test_dqn.py	100.00% <100.00%> (ø)
rlberry/stats/evaluation.py	87.01% <100.00%> (-0.17%)	⬇️
... and 13 more

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 1e36b18...806a726. Read the comment docs.