[Feature Request] ACME Integration

[Trinkle23897]

https://github.com/deepmind/acme

Road Map:

@TianyiSun316

• Go through ACME codebase and integrate vector_env to the available algorithms;
• Write Atari examples;
• Check Atari performance: Pong and Breakout;
• Submit PR;

---

@LeoGuo98

• Do some experiments with sample efficiency (actually you can try out with different libraries, either ACME, tianshou, or sb3, this doesn't depend on the previous item)

Resources:

tianshou: #51
stable-baselines3: #39
cleanrl: #48 #53

cc @zhongwen

[zhongwen]

Use Acme JAX agents intead of the TF agents

References:
https://github.com/deepmind/acme/tree/master/acme/agents/jax

Example scripts, you can just replace it with a dqn agent.

https://github.com/deepmind/acme/blob/master/examples/atari/run_impala.py

https://github.com/deepmind/acme/blob/master/examples/atari/run_r2d2.py

[TianyiSun316]

Use Acme JAX agents intead of the TF agents

References: https://github.com/deepmind/acme/tree/master/acme/agents/jax

Example scripts, you can just replace it with a dqn agent.

https://github.com/deepmind/acme/blob/master/examples/atari/run_impala.py

https://github.com/deepmind/acme/blob/master/examples/atari/run_r2d2.py

I tried but seems like JAX DQN network is not usable. Current version is based on TensorFlow. Maybe we can discuss this offline?

[zhongwen]

Sure!

[zhongwen]

Let's use R2D2 example first, since Tianyi has significant difficutlites and spent little time to understand the acme codebase.

[zhongwen]

Steps:

• Implement a new EnvironmentLoop to integrate with EnvPool;
• Impelement a new adder to integrate with batched inputs;
• Integrate with the R2D2 example
• Experiments

[zhongwen]

How is the progress? You have only 1 day left for the task.

[TianyiSun316]

Let's use R2D2 example first, since Tianyi has significant difficutlites and spent little time to understand the acme codebase.

I was eager to know your solution of using JAX DQN to write the example when we discuss offline and there are several things I forgot to mention during the offline discussion:

1. I have already finished step 1 in your "Steps", you will find it in my PR Add acme JAX R2D2 example #104
2. Using JAX DQN is more complex than the example scripts you mentioned above. There is no pre-written DQN network that we can use directly. We may need to write functions like make_atari_networks to get a usable DQN network, while the example scripts simply called the API that has already been implemented.
3. Implementing R2D2 is also more complex than the example scripts since the R2D2 agent expects the observation to use the OAR struct. Our EnvWrapper needs to adapt to it accordingly.

I have finished the example using JAX R2D2 agents. Please review the code and let me know if there are any modifications needed.