From 2ad5c107b88b14dca6d94c740beecae5d37e063a Mon Sep 17 00:00:00 2001
From: kngwyu <yuji.kngw.80s.revive@gmail.com>
Date: Fri, 23 Aug 2024 14:33:25 +0900
Subject: [PATCH] Agents dies as soon as energy exhausted

---
 experiments/cf_simple.py | 8 +++++++-
 src/emevo/env.py         | 2 +-
 tests/test_status.py     | 2 +-
 3 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/experiments/cf_simple.py b/experiments/cf_simple.py
index 7344ed25..1fb14155 100644
--- a/experiments/cf_simple.py
+++ b/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(
diff --git a/src/emevo/env.py b/src/emevo/env.py
index eefb624a..f0724397 100644
--- a/src/emevo/env.py
+++ b/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:
diff --git a/tests/test_status.py b/tests/test_status.py
index 7f0963a5..57321c8e 100644
--- a/tests/test_status.py
+++ b/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)