Fix issue #3. (#5)

* Fix issue #3. * Use the field from flax struct. * Remove unused import.
dunnolab · Feb 1, 2024 · e574386 · e574386
1 parent abf181c
commit e574386
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/xminigrid/types.py b/src/xminigrid/types.py
@@ -18,9 +18,9 @@ class RuleSet(struct.PyTreeNode):
 
 
 class AgentState(struct.PyTreeNode):
-    position: jax.Array = jnp.asarray((0, 0))
-    direction: jax.Array = jnp.asarray(0)
-    pocket: jax.Array = TILES_REGISTRY[Tiles.EMPTY, Colors.EMPTY]
+    position: jax.Array = struct.field(default_factory=lambda: jnp.asarray((0, 0)))
+    direction: jax.Array = struct.field(default_factory=lambda: jnp.asarray(0))
+    pocket: jax.Array = struct.field(default_factory=lambda: TILES_REGISTRY[Tiles.EMPTY, Colors.EMPTY])
 
 
 class EnvCarry(struct.PyTreeNode):