[experimental] Add error message to auto_reset (#1102)

Author: sotetsuk

pgx/experimental/wrappers.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 import jax
 import jax.numpy as jnp
 
@@ -29,7 +31,16 @@ def auto_reset(step_fn, init_fn):
     2. Performance
     """
 
-    def wrapped_step_fn(state: State, action: Array, key: PRNGKey):
+    def wrapped_step_fn(
+        state: State, action: Array, key: Optional[PRNGKey] = None
+    ):
+        assert key is not None, (
+            "v2.0.0 changes the signature of auto reset. Please specify PRNGKey at the third argument:\n\n"
+            "  * <  v2.0.0: step_fn(state, action)\n"
+            "  * >= v2.0.0: step_fn(state, action, key)\n\n"
+            "Note that codes under pgx.experimental are subject to change without notice."
+        )
+
         key1, key2 = jax.random.split(key)
         state = jax.lax.cond(
             (state.terminated | state.truncated),