Fix some names leftover from brax example

Signed-off-by: Fabrice Normandin <[email protected]>

benchmarks/purejaxrl/benchfile.py

@@ -1,9 +1,9 @@
 from milabench.pack import Package
 
 
-class BraxBenchmark(Package):
+class PureJaxRLBenchmark(Package):
     base_requirements = "requirements.in"
     main_script = "main.py"
 
 
-__pack__ = BraxBenchmark
+__pack__ = PureJaxRLBenchmark()

benchmarks/purejaxrl/main.py

@@ -1,4 +1,10 @@
-"""PureJaxRL"""
+"""PureJaxRL example based on https://github.com/luchris429/purejaxrl/blob/main/examples/walkthrough.ipynb
+
+TODOs:
+- This is a simple example on CartPole, it would be preferable to do something more serious (and for longer). 
+- Should probably try something like https://github.com/luchris429/purejaxrl/blob/main/purejaxrl/experimental/s5/ppo_s5.py
+
+"""
 import dataclasses
 import os
 

benchmarks/purejaxrl/voirfile.py

@@ -1,8 +1,9 @@
 from dataclasses import dataclass
 
 from voir import configurable
+from voir.instruments import dash, early_stop, log, rate
 from voir.phase import StopProgram
-from voir.instruments import dash, early_stop, gpu_monitor, log, rate
+
 from benchmate.monitor import monitor_monogpu
 
 
@@ -46,4 +47,4 @@ def instrument_main(ov, options: Config):
     try:
         yield ov.phases.run_script
     except StopProgram:
-        print("early stopped")
+        print("early stopped")