Agents dies as soon as energy exhausted

experiments/cf_simple.py

@@ -120,7 +120,13 @@ def step_rollout(
         )
         # Birth and death
         death_prob = hazard_fn(state_t1.status.age, state_t1.status.energy)
-        dead = jax.random.bernoulli(hazard_key, p=death_prob)
+        dead_nonzero = jax.random.bernoulli(hazard_key, p=death_prob)
+        dead = jnp.where(
+            # If the agent's energy is lower than 0, it should immediately die
+            state_t1.status.energy < 0.0,
+            jnp.ones_like(dead_nonzero),
+            dead_nonzero,
+        )
         state_t1d = env.deactivate(state_t1, dead)
         birth_prob = birth_fn(state_t1d.status.age, state_t1d.status.energy)
         possible_parents = jnp.logical_and(

src/emevo/env.py

@@ -49,7 +49,7 @@ def deactivate(self, flag: jax.Array) -> Self:
     def update(self, energy_delta: jax.Array, capacity: float | None = 100.0) -> Self:
         """Update energy."""
         energy = self.energy + energy_delta
-        return replace(self, energy=jnp.clip(energy, a_min=0.0, a_max=capacity))
+        return replace(self, energy=jnp.clip(energy, a_max=capacity))
 
 
 def init_status(max_n: int, init_energy: float) -> Status:

tests/test_status.py

@@ -15,7 +15,7 @@ def test_status_clipping(n: int, capacity: float) -> None:
         assert jnp.all(status.energy >= 0.0)
         assert jnp.all(status.energy <= capacity)
 
-    for _ in range(300):
+    for _ in range(200):
         status.update(energy_delta=jnp.ones(n) * -1.0, capacity=capacity)
         assert jnp.all(status.energy >= 0.0)
         assert jnp.all(status.energy <= capacity)