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?