go.py: state.env_id throws exception when running tutorial example

[DuaneNielsen]

In go.py

    def env_id(self) -> core.EnvId:
        try:
            size = int(self._size.item())
        except TypeError:
            size = int(self._size[0].item())
        return f"go_{size}x{size}"  # type: ignore

I"m new to the code, but it looks like the intent here was to pick the first state on the array if the state is vectorized.

On my system, if the state is vectorized.. this throws a ValueError, not a Type Error.. like below

import jax
import pgx
from pgx.experimental import act_randomly

batch_size = 9

env = pgx.make("go_9x9")

init = jax.jit(jax.vmap(env.init))
step = jax.jit(jax.vmap(env.step))
act_randomly = jax.jit(act_randomly)

key = jax.random.PRNGKey(42)
key, subkey = jax.random.split(key)
keys = jax.random.split(subkey, batch_size)

state = init(keys)
state.env_id

Traceback (most recent call last): File "/home/duane/PycharmProjects/mctx_bench/.venv/lib/python3.10/site-packages/pgx/go.py", line 56, in env_id size = int(self._size.item()) File "/home/duane/PycharmProjects/mctx_bench/.venv/lib/python3.10/site-packages/jax/_src/numpy/array_methods.py", line 81, in _item return arr.item(*args) ValueError: can only convert an array of size 1 to a Python scalar

Not sure what the fix is.. maybe

    def env_id(self) -> core.EnvId:
        if len(self._size) == 0:
            size = int(self._size.item())
        else:
            size = int(self._size[0].item())
        return f"go_{size}x{size}"  # type: ignore

[DuaneNielsen]

I think this maybe..

    @property
    def env_id(self) -> core.EnvId:
        if self._size.ndim == 0:
            size = int(self._size.item())
        else:
            size = int(self._size[0].item())
        return f"go_{size}x{size}"  # type: ignore

[DuaneNielsen]

Tested the above.. seems to work.

[sotetsuk]

Hi, thank you for reporting a bug! 🙏 We released a new version which should resolve it. Please let me know if you find another bug!