Fix rendering of the SlimeVolley by extracting state

[hota911]

What

This change fixes the error of SlimeVolley.render fucntion.

• Running train_slimevolley.py fails due to the following error: TypeError: Only scalar arrays can be converted to Python scalars; got arr.ndim=1
  • State is vectorized by jax.vmap, but display() functions assume the original shape of the task.
• In this PR, we use jax.tree.map to extract the state of one task.

Error log

$ uv run python examples/train_slimevolley.py --max-iter=1 --n-repeats=1
SlimeVolley: 2024-12-26 15:18:24,981 [INFO] EvoJAX SlimeVolley
SlimeVolley: 2024-12-26 15:18:24,982 [INFO] ==============================
INFO:2024-12-26 15:18:25,014:jax._src.xla_bridge:927: Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
jax._src.xla_bridge: 2024-12-26 15:18:25,014 [INFO] Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
INFO:2024-12-26 15:18:25,016:jax._src.xla_bridge:927: Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file), '/usr/local/lib/libtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache)
jax._src.xla_bridge: 2024-12-26 15:18:25,016 [INFO] Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file), '/usr/local/lib/libtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache)
MLPPolicy: 2024-12-26 15:18:26,008 [INFO] MLPPolicy.num_params = 323
(64_w,128)-aCMA-ES (mu_w=34.2,w_1=6%) in dimension 323 (seed=123, Thu Dec 26 15:18:40 2024)
SlimeVolley: 2024-12-26 15:18:40,721 [INFO] use_for_loop=False
SlimeVolley: 2024-12-26 15:18:40,746 [INFO] Start to train for 1 iterations.
SlimeVolley: 2024-12-26 15:18:43,622 [INFO] [TEST] Iter=1, #tests=100, max=-3.0000, avg=-4.8200, min=-5.0000, std=0.4556
SlimeVolley: 2024-12-26 15:18:43,643 [INFO] Training done, best_score=-4.8200
SlimeVolley: 2024-12-26 15:18:43,646 [INFO] Loaded model parameters from ./log/slimevolley.
SlimeVolley: 2024-12-26 15:18:43,646 [INFO] Start to test the parameters.
SlimeVolley: 2024-12-26 15:18:43,924 [INFO] [TEST] #tests=100, max=-2.0000, avg=-4.8200, min=-5.0000, std=0.4556
Traceback (most recent call last):
  File "/Users/hota/workspace/hota911/evojax/examples/train_slimevolley.py", line 154, in <module>
    main(configs)
  File "/Users/hota/workspace/hota911/evojax/examples/train_slimevolley.py", line 142, in main
    screens.append(SlimeVolley.render(task_state))
  File "/Users/hota/workspace/hota911/evojax/evojax/task/slimevolley.py", line 902, in render
    canvas = game.display()
  File "/Users/hota/workspace/hota911/evojax/evojax/task/slimevolley.py", line 759, in display
    canvas = self.agent_left.display(canvas, self.ball.p.x, self.ball.p.y)
  File "/Users/hota/workspace/hota911/evojax/evojax/task/slimevolley.py", line 594, in display
    bx = float(ball_x)
  File "/Users/hota/workspace/hota911/evojax/.venv/lib/python3.10/site-packages/jax/_src/array.py", line 298, in __float__
    core.check_scalar_conversion(self)
  File "/Users/hota/workspace/hota911/evojax/.venv/lib/python3.10/site-packages/jax/_src/core.py", line 641, in check_scalar_conversion
    raise TypeError("Only scalar arrays can be converted to Python scalars; "
TypeError: Only scalar arrays can be converted to Python scalars; got arr.ndim=1

QA

With this fix:

$ python examples/train_slimevolley.py --max-iter=1 --n-repeats=1
SlimeVolley: 2024-12-26 15:20:05,806 [INFO] EvoJAX SlimeVolley
SlimeVolley: 2024-12-26 15:20:05,806 [INFO] ==============================
INFO:2024-12-26 15:20:05,823:jax._src.xla_bridge:927: Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
jax._src.xla_bridge: 2024-12-26 15:20:05,823 [INFO] Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
INFO:2024-12-26 15:20:05,824:jax._src.xla_bridge:927: Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file), '/usr/local/lib/libtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache)
jax._src.xla_bridge: 2024-12-26 15:20:05,824 [INFO] Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: dlopen(libtpu.so, 0x0001): tried: 'libtpu.so' (no such file), '/System/Volumes/Preboot/Cryptexes/OSlibtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache), 'libtpu.so' (no such file), '/usr/local/lib/libtpu.so' (no such file), '/usr/lib/libtpu.so' (no such file, not in dyld cache)
MLPPolicy: 2024-12-26 15:20:06,617 [INFO] MLPPolicy.num_params = 323
(64_w,128)-aCMA-ES (mu_w=34.2,w_1=6%) in dimension 323 (seed=123, Thu Dec 26 15:20:07 2024)
SlimeVolley: 2024-12-26 15:20:07,250 [INFO] use_for_loop=False
SlimeVolley: 2024-12-26 15:20:07,273 [INFO] Start to train for 1 iterations.
SlimeVolley: 2024-12-26 15:20:10,172 [INFO] [TEST] Iter=1, #tests=100, max=-3.0000, avg=-4.8200, min=-5.0000, std=0.4556
SlimeVolley: 2024-12-26 15:20:10,196 [INFO] Training done, best_score=-4.8200
SlimeVolley: 2024-12-26 15:20:10,197 [INFO] Loaded model parameters from ./log/slimevolley.
SlimeVolley: 2024-12-26 15:20:10,197 [INFO] Start to test the parameters.
SlimeVolley: 2024-12-26 15:20:10,473 [INFO] [TEST] #tests=100, max=-2.0000, avg=-4.8200, min=-5.0000, std=0.4556
SlimeVolley: 2024-12-26 15:20:31,551 [INFO] GIF saved to ./log/slimevolley/slimevolley.gif.

Ref

• The error was reported by other user: Type Conflict Between JAX NumPy and Python's float() in EvoJAX SlimeVolley #81