train error - Githubissues

Hello. I tried experimenting with the rules so I could train yachzotron for a slightly different game. Unfortunately, call python3 yahtzotron/cli.py train -o custom.pkl --ruleset yatzy_modified fails with

  0%|                                                             | 1/20000 [00:01<7:01:30,  1.26s/it, actor_loss=8.04, critic_loss=77, entropy_loss=-0.000279, loss=85.1, score=173]
Traceback (most recent call last):
  File "/home/user/storage/lab/yahtzotron/yahtzotron/cli.py", line 171, in <module>
    cli()
  File "/home/user/storage/lab/yahtzotron/venv/lib/python3.12/site-packages/click/core.py", line 1157, in __call__
    return self.main(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/storage/lab/yahtzotron/venv/lib/python3.12/site-packages/click/core.py", line 1078, in main
    rv = self.invoke(ctx)
         ^^^^^^^^^^^^^^^^
  File "/home/user/storage/lab/yahtzotron/venv/lib/python3.12/site-packages/click/core.py", line 1688, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/storage/lab/yahtzotron/venv/lib/python3.12/site-packages/click/core.py", line 1434, in invoke
    return ctx.invoke(self.callback, **ctx.params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/storage/lab/yahtzotron/venv/lib/python3.12/site-packages/click/core.py", line 783, in invoke
    return __callback(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/storage/lab/yahtzotron/yahtzotron/cli.py", line 70, in train
    yzt = train_a2c(yzt, num_epochs=20_000, pretraining=True)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/storage/lab/yahtzotron/venv/lib/python3.12/site-packages/yahtzotron/training.py", line 197, in train_a2c
    weights, opt_state = sgd_step(
                         ^^^^^^^^^
  File "/home/user/storage/lab/yahtzotron/venv/lib/python3.12/site-packages/yahtzotron/training.py", line 90, in sgd_step
    updates, opt_state = optimizer.update(gradients, opt_state)
                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/storage/lab/yahtzotron/venv/lib/python3.12/site-packages/optax/transforms/_accumulation.py", line 380, in update
    new_updates, new_state = lax.cond(
                             ^^^^^^^^^
  File "/home/user/storage/lab/yahtzotron/venv/lib/python3.12/site-packages/optax/transforms/_accumulation.py", line 336, in _do_update
    acc_grads = jtu.tree_map(
                ^^^^^^^^^^^^^
ValueError: Custom node type mismatch: expected type: <class 'haiku._src.data_structures.FlatMap'>, value: {'linear': {'b': Traced<ShapedArray(float32[128])>with<DynamicJaxprTrace(level=2/0)>, 'w': Traced<ShapedArray(float32[25,128])>with<DynamicJaxprTrace(level=2/0)>}, 'linear_1': {'b': Traced<ShapedArray(float32[256])>with<DynamicJaxprTrace(level=2/0)>, 'w': Traced<ShapedArray(float32[128,256])>with<DynamicJaxprTrace(level=2/0)>}, 'linear_2': {'b': Traced<ShapedArray(float32[128])>with<DynamicJaxprTrace(level=2/0)>, 'w': Traced<ShapedArray(float32[256,128])>with<DynamicJaxprTrace(level=2/0)>}, 'linear_3': {'b': Traced<ShapedArray(float32[1])>with<DynamicJaxprTrace(level=2/0)>, 'w': Traced<ShapedArray(float32[128,1])>with<DynamicJaxprTrace(level=2/0)>}, 'linear_4': {'b': Traced<ShapedArray(float32[32])>with<DynamicJaxprTrace(level=2/0)>, 'w': Traced<ShapedArray(float32[128,32])>with<DynamicJaxprTrace(level=2/0)>}, 'linear_5': {'b': Traced<ShapedArray(float32[15])>with<DynamicJaxprTrace(level=2/0)>, 'w': Traced<ShapedArray(float32[128,15])>with<DynamicJaxprTrace(level=2/0)>}}.
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

Same goes for unmodified rules. My guess is the vague setup.py file, which does not specify the versions of the packages, and newer versions lost compatibility with your code. Is there any chance to restore the knowledge of the environment in which yachzotron was created and share it?

dionhaefner / yahtzotron

train error #3