Implementing an agent from the literature

ltc/agents/__init__.py

@@ -1,3 +1,8 @@
 from ltc.agents.dcf import DCF
-from ltc.agents.deep import BayesianDDQN
+from ltc.agents.deep import BayesianDDQN, DDQN
+from ltc.agents.dlma import DLMANetwork
 from ltc.agents.drl import QNetwork, QNetworkDropout, StochasticVariationalNetwork
+from ltc.agents.eb_aloha import EBALOHA
+from ltc.agents.fw_aloha import FWALOHA
+from ltc.agents.q_aloha import QALOHA
+from ltc.agents.tdma import TDMA

ltc/agents/dlma.py

@@ -0,0 +1,89 @@
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+import flax.linen as nn
+
+from ltc.agents.drl import add_batch_dim
+
+
+@jax.vmap
+@jax.vmap
+def prepare_observation(obs):
+    def prepare_action(obs):
+        return jax.nn.one_hot(obs[1], 2)
+
+    def prepare_state(obs):
+        """
+        **WARNING**: This function assumes that there are only two nodes in the network (1 DLMA and 1 legacy)!
+
+        1. ch = -1, act = 0  (collision)
+        2. ch = 1, act = 0   (DLMA transmits)
+        3. ch = 1, act = 1   (legacy transmits)
+        4. ch = 0, act = 1   (idle)
+
+        s = max(ch + 2 * act, 0)
+
+        1. s = max(-1 + 0, 0) = 0  (collision)
+        2. s = max(1 + 0, 0) = 1   (DLMA transmits)
+        3. s = max(1 + 2, 0) = 3   (legacy transmits)
+        4. s = max(0 + 2, 0) = 2   (idle)
+        """
+
+        state = obs[:-1]  # remove buffer state
+        state = state[0] + 2 * state[1]
+        state = jnp.maximum(state, 0)
+        state = jax.nn.one_hot(state, 4)
+        return state
+
+    def prepare_rewards(obs):
+        """
+        **WARNING**: This function assumes that there are only two nodes in the network (1 DLMA and 1 legacy)!
+
+        1. ch = -1, act = 0  (collision)        -> DLMA reward = 0, legacy reward = 0
+        2. ch = 1, act = 0   (DLMA transmits)   -> DLMA reward = 1, legacy reward = 0
+        3. ch = 1, act = 1   (legacy transmits) -> DLMA reward = 0, legacy reward = 1
+        4. ch = 0, act = 1   (idle)             -> DLMA reward = 0, legacy reward = 0
+        """
+
+        state = obs[:-1]  # remove buffer state
+        no_tx = jnp.maximum(state[0], 0)
+        dlma_tx = no_tx * (1 - state[1])
+        legacy_tx = no_tx * state[1]
+        reward = jnp.array([dlma_tx, legacy_tx])
+        return reward
+
+    action = prepare_action(obs)
+    state = prepare_state(obs)
+    reward = prepare_rewards(obs)
+    obs = jnp.concatenate([action, state, reward], axis=0)
+    return obs
+
+
+class DLMANetwork(nn.Module):
+    num_actions: int = 2
+
+    @nn.compact
+    def __call__(self, s):
+        dense = partial(nn.Dense, kernel_init=nn.initializers.he_normal())
+
+        s = add_batch_dim(s)
+        s = prepare_observation(s)
+
+        b, *_ = s.shape
+        x = s.reshape(b, -1)
+
+        for _ in range(2):
+            x = dense(64)(x)
+            x = nn.relu(x)
+
+        for _ in range(2):
+            x_res = x
+            x = dense(64)(x)
+            x = nn.relu(x)
+            x = dense(64)(x)
+            x = nn.relu(x)
+            x = x + x_res
+
+        x = dense(self.num_actions)(x)
+        return x

ltc/agents/eb_aloha.py

@@ -0,0 +1,57 @@
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+from chex import dataclass
+from reinforced_lib.agents import BaseAgent, AgentState
+
+from ltc.sim.constants import Actions
+
+
+@dataclass
+class EBALOHAState(AgentState):
+    backoff: int
+    ret_c: int
+
+
+class EBALOHA(BaseAgent):
+    def __init__(self, window_size: int, max_backoff: int):
+        self.window_size = window_size
+        self.max_backoff = max_backoff
+        self.init = jax.jit(partial(self.init, window_size=window_size))
+        self.update = jax.jit(partial(self.update, window_size=window_size, max_backoff=max_backoff))
+        self.sample = jax.jit(self.sample)
+
+    @staticmethod
+    def init(key, window_size: int):
+        ret_c = 0
+        backoff = jax.random.randint(key, shape=(), minval=0, maxval=window_size)
+        return EBALOHAState(backoff=backoff, ret_c=ret_c)
+
+    @staticmethod
+    def update(state, key, env_state, action, reward, terminal, window_size: int, max_backoff: int):
+        collision = (action == Actions.TX.value) & (reward == 0)
+        success = (action == Actions.TX.value) & (reward > 0)
+
+        new_ret_c = jnp.where(
+            collision,
+            jnp.minimum(state.ret_c + 1, max_backoff),
+            jnp.where(success, 0, state.ret_c)
+        )
+        window = window_size * (2 ** new_ret_c)
+
+        reset_value = jax.random.randint(key, shape=(), minval=0, maxval=window)
+        new_backoff = state.backoff - 1
+        new_backoff = jnp.where(new_backoff < 0, reset_value, new_backoff)
+
+        return EBALOHAState(backoff=new_backoff, ret_c=new_ret_c)
+
+    @staticmethod
+    def sample(state, key, env_state):
+        *_, buffer = env_state[-1]
+
+        return jnp.where(
+            (buffer > 0) & (state.backoff == 0),
+            Actions.TX.value,
+            Actions.IDLE.value
+        )

ltc/agents/fw_aloha.py

@@ -0,0 +1,43 @@
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+from chex import dataclass
+from reinforced_lib.agents import BaseAgent, AgentState
+
+from ltc.sim.constants import Actions
+
+
+@dataclass
+class FWALOHAState(AgentState):
+    backoff: int
+
+
+class FWALOHA(BaseAgent):
+    def __init__(self, window_size: int):
+        self.window_size = window_size
+        self.init = jax.jit(partial(self.init, window_size=window_size))
+        self.update = jax.jit(partial(self.update, window_size=window_size))
+        self.sample = jax.jit(self.sample)
+
+    @staticmethod
+    def init(key, window_size: int):
+        backoff = jax.random.randint(key, shape=(), minval=0, maxval=window_size)
+        return FWALOHAState(backoff=backoff)
+
+    @staticmethod
+    def update(state, key, env_state, action, reward, terminal, window_size: int):
+        reset_value = jax.random.randint(key, shape=(), minval=0, maxval=window_size)
+        new_backoff = state.backoff - 1
+        new_backoff = jnp.where(new_backoff < 0, reset_value, new_backoff)
+        return FWALOHAState(backoff=new_backoff)
+
+    @staticmethod
+    def sample(state, key, env_state):
+        *_, buffer = env_state[-1]
+
+        return jnp.where(
+            (buffer > 0) & (state.backoff == 0),
+            Actions.TX.value,
+            Actions.IDLE.value
+        )

ltc/agents/q_aloha.py

@@ -0,0 +1,42 @@
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+from chex import dataclass
+from reinforced_lib.agents import BaseAgent, AgentState
+
+from ltc.sim.constants import Actions
+
+
+@dataclass
+class QALOHAState(AgentState):
+    pass
+
+
+class QALOHA(BaseAgent):
+    def __init__(self, q):
+        self.init = jax.jit(self.init)
+        self.update = jax.jit(self.update)
+        self.sample = jax.jit(partial(self.sample, q=q))
+
+    @staticmethod
+    def init(key):
+        return QALOHAState()
+
+    @staticmethod
+    def update(state, key, env_state, action, reward, terminal):
+        return state
+
+    @staticmethod
+    def sample(state, key, env_state, q):
+        *_, buffer = env_state[-1]
+
+        return jnp.where(
+            buffer == 0,
+            Actions.IDLE.value,
+            jnp.where(
+                jax.random.uniform(key) < q,
+                Actions.TX.value,
+                Actions.IDLE.value
+            )
+        )

ltc/agents/tdma.py

@@ -0,0 +1,44 @@
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+from chex import dataclass
+from reinforced_lib.agents import BaseAgent, AgentState
+
+from ltc.sim.constants import Actions
+
+
+@dataclass
+class TDMAState(AgentState):
+    counter: int
+    tdma_slots: jnp.ndarray
+
+
+class TDMA(BaseAgent):
+    def __init__(self, state_size: int, assigned_slots: int):
+        self.state_size = state_size
+        self.assigned_slots = assigned_slots
+        self.init = jax.jit(partial(self.init, state_size=state_size, assigned_slots=assigned_slots))
+        self.update = jax.jit(self.update)
+        self.sample = jax.jit(self.sample)
+
+    @staticmethod
+    def init(key, state_size: int, assigned_slots: int):
+        tdma_slots = jnp.zeros(state_size, dtype=int).at[:assigned_slots].set(1)
+        tdma_slots = jax.random.permutation(key, tdma_slots)
+        return TDMAState(counter=0, tdma_slots=tdma_slots)
+
+    @staticmethod
+    def update(state, key, env_state, action, reward, terminal):
+        new_count = (state.counter + 1) % state.tdma_slots.shape[0]
+        return TDMAState(counter=new_count, tdma_slots=state.tdma_slots)
+
+    @staticmethod
+    def sample(state, key, env_state):
+        *_, buffer = env_state[-1]
+
+        return jnp.where(
+            (buffer > 0) & (state.tdma_slots[state.counter] == 1),
+            Actions.TX.value,
+            Actions.IDLE.value
+        )