diff --git a/README.md b/README.md index 97296246..483fddb2 100644 --- a/README.md +++ b/README.md @@ -73,9 +73,9 @@ Stoix natively supports logging to json files which adhere to the standard sugge As the code in Stoix (at the time of creation) was in essense a port of [Mava][mava], for further speed comparisons we point to their repo. Additionally, we refer to the PureJaxRL blog post [here](https://chrislu.page/blog/meta-disco/) where the speed benefits of end-to-end JAX systems are discussed. Below we provide some plots illustrating that Stoix performs equally to that of [PureJaxRL][purejaxrl] but with the added benefit of the code being already set up for `pmap` distribution over devices as well as the other features provided (algorithm implementations, logging, config system, etc). - -ppo dqn - +

+ppo dqn +

I've also included a plot of the training time for 5e5 steps of PPO as one scales the number of environments. PureJaxRL does not pmap and thus runs on a single a device.