diff --git a/src/xminigrid/types.py b/src/xminigrid/types.py
index 326437b..7a69c0b 100644
--- a/src/xminigrid/types.py
+++ b/src/xminigrid/types.py
@@ -17,9 +17,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):