V2

.gitignore

@@ -169,3 +169,6 @@ cython_debug/
 
 # PyPI configuration file
 .pypirc
+
+scripts/
+plots/

ltc/agents/dcf.py

@@ -3,7 +3,7 @@
 from chex import dataclass
 from reinforced_lib.agents import BaseAgent, AgentState
 
-from ltc.sim.constants import Actions
+from ltc.sim.constants import Actions, ObsIdx
 
 
 @dataclass
@@ -44,22 +44,28 @@ def double_cw():
             backoff = jax.random.randint(key, (), 0, state.cw)
             return DCFState(cw=cw, backoff=backoff)
 
-        buffer, channel, ret_c, _, _ = env_state[-1]
+        buffer_count = env_state[-1][ObsIdx.BUFFER_PACKET_COUNT]
+        channel_busy = env_state[-1][ObsIdx.CHANNEL_LAST_CS_BUSY]
+        ret_c = env_state[-1][ObsIdx.STATUS_RETRY_COUNTER]
+
+        is_tx = action == Actions.TX.value
+        is_success = jax.lax.bitwise_and(is_tx, jax.lax.bitwise_or(reward > 0, ret_c == 0))
+        is_failure = jax.lax.bitwise_and(is_tx, reward < 0)
 
         return jax.lax.cond(
-            buffer == 0,
+            buffer_count == 0,
             reset,
             lambda: jax.lax.cond(
-                jax.lax.bitwise_and(state.backoff > 0, channel == 0),
+                jax.lax.bitwise_and(state.backoff > 0, channel_busy == 0),
                 countdown,
                 lambda: jax.lax.cond(
-                    jax.lax.bitwise_and(state.backoff > 0, channel != 0),
+                    jax.lax.bitwise_and(state.backoff > 0, channel_busy != 0),
                     freeze,
                     lambda: jax.lax.cond(
-                        jax.lax.bitwise_or(reward > 0, ret_c == 0),
+                        is_success,
                         reset,
                         lambda: jax.lax.cond(
-                            reward < 0,
+                            is_failure,
                             double_cw,
                             freeze
                         )
@@ -70,13 +76,14 @@ def double_cw():
 
     @staticmethod
     def sample(state, key, env_state):
-        buffer, channel, _, _, _ = env_state[-1]
+        buffer_count = env_state[-1][ObsIdx.BUFFER_PACKET_COUNT]
+        channel_busy = env_state[-1][ObsIdx.CHANNEL_LAST_CS_BUSY]
 
         return jnp.where(
-            buffer == 0,
+            buffer_count == 0,
             Actions.IDLE.value,
             jnp.where(
-                jax.lax.bitwise_and(state.backoff == 0, channel == 0),
+                jax.lax.bitwise_and(state.backoff == 0, channel_busy == 0),
                 Actions.TX.value,
                 Actions.CS.value
             )