Addition of ACK functionality

ltc/run.py

@@ -1,4 +1,5 @@
 import os
+
 os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = 'false'
 
 from dataclasses import replace
@@ -13,7 +14,7 @@
 
 from ltc.agents import BayesianDDQN, DCF, QNetwork, QNetworkDropout, StochasticVariationalNetwork
 from ltc.sim import InitialStateConf, cox_traffic, process_output, simulate
-from ltc.sim.constants import INITIAL_CAPACITY, Actions
+from ltc.sim.constants import INITIAL_CAPACITY, Actions, ObservationIndex
 from ltc.utils.scan_states import Carry, Output
 from ltc.utils.plots import plot_all, plot_first
 
@@ -51,7 +52,7 @@ def init_traffic(traffic, key, n):
 
 def rl_step(drl_step, dcf_step, traffic_step, n, n_drl, n_bins=50):
     def rl_step_fn(c, _):
-        key, drl_keys, dcf_keys, traffic_key = jax.random.split(c.key, 4)
+        key, drl_keys, dcf_keys, traffic_key, simulation_key = jax.random.split(c.key, 5)
         drl_keys = jax.random.split(drl_keys, n_drl)
         dcf_keys = jax.random.split(dcf_keys, n - n_drl)
         traffic_keys = jax.random.split(traffic_key, n)
@@ -61,8 +62,10 @@ def rl_step_fn(c, _):
         actions = jnp.concatenate([drl_actions, dcf_actions])
 
         traffic_states, new_frames = traffic_step(c.traffic_states, traffic_keys)
-        buffer_states, channel_state = simulate(c.buffer_states, new_frames, actions)
-        obs, rewards, powers = process_output(c.buffer_states, buffer_states, c.power_states, channel_state, c.obs, actions, c.terminals)
+        buffer_states, channel_state, transmission_history = simulate(simulation_key, c.buffer_states, new_frames,
+                                                                      actions, c.transmission_history, n)
+        obs, rewards, powers = process_output(c.buffer_states, buffer_states, c.power_states, channel_state, c.obs,
+                                              actions, c.terminals, transmission_history)
         terminals = jnp.logical_or(c.terminals, powers < 0)
 
         params = c.drl_states.params['model'] if 'model' in c.drl_states.params else c.drl_states.params
@@ -72,7 +75,7 @@ def rl_step_fn(c, _):
         hist, bin_edges = jax.vmap(jnp.histogram, in_axes=(0, None))(flat_params, n_bins)
 
         c = Carry(
-            drl_states, dcf_states, traffic_states, buffer_states, powers,
+            drl_states, dcf_states, traffic_states, buffer_states, powers, transmission_history,
             channel_state, key, obs, actions, rewards, terminals
         )
         o = Output(
@@ -95,7 +98,10 @@ def rl_step_fn(c, _):
     buffer_states = jnp.zeros(n, dtype=int)
     power_states = jnp.full(n, INITIAL_CAPACITY, dtype=int)
     channel_state = 0
-    obs = jnp.zeros((n, window_size, 5), dtype=int).at[:, -1].set(INITIAL_CAPACITY)
+    obs = jnp.zeros((n, window_size, 6), dtype=int).at[:, ObservationIndex.POWER.value].set(INITIAL_CAPACITY)
+    id_values = jnp.arange(0, n).reshape(n, 1)
+    obs = obs.at[:, :, ObservationIndex.ID.value].set(id_values)
+    transmission_history = jnp.zeros((window_size + 1, 2))
     rewards = jnp.zeros(n)
     terminals = jnp.full(n, False, dtype=bool)
 
@@ -127,7 +133,7 @@ def rl_step_fn(c, _):
 
     rl_step_fn = jax.jit(rl_step(drl_step, dcf_step, traffic_step, n, n_drl))
     init_carry = Carry(
-        drl_states, dcf_states, traffic_states, buffer_states, power_states,
+        drl_states, dcf_states, traffic_states, buffer_states, power_states, transmission_history,
         channel_state, key, obs, actions, rewards, terminals
     )
     all_outputs = []

ltc/sim/constants.py

@@ -2,12 +2,10 @@
 
 
 class Actions(Enum):
-    """
-    Actions for the agent
-    """
     TX = 0
     CS = 1
     IDLE = 2
+    ACK = 3
 
 
 """
@@ -21,10 +19,12 @@ class Actions(Enum):
 Rewards and penalties
 """
 TX_REWARD = 1.0
+ACK_REWARD = 1.0
 EMPTY_BUFFER_REWARD = 0.5
 NO_TX_REWARD = 0.0
 NO_TX_PENALTY = -1.0
 EMPTY_TX_PENALTY = -0.5
+EMPTY_ACK_PENALTY = -0.5
 COLLISION_PENALTY = -1.0
 MAX_RETRANSMISSION_PENALTY = -1.0
 
@@ -33,6 +33,7 @@ class Actions(Enum):
 """
 INITIAL_CAPACITY = 10000
 TX_CONSUMPTION = 5
+ACK_CONSUMPTION = 2
 CS_CONSUMPTION = 1
 IDLE_CONSUMPTION = 0
 
@@ -41,3 +42,22 @@ class Actions(Enum):
 https://ieeexplore.ieee.org/document/8930559
 """
 TAU = 5.484 * 1e-3
+
+
+class AckState(Enum):
+    SENT = 1
+    NOT_SENT = 0
+
+
+class ObservationIndex(Enum):
+    BUFFER = 0
+    CHANNEL = 1
+    RET_C = 2
+    NO_TX = 3
+    POWER = 4
+    ID = 5
+
+
+class TransmissionIndex(Enum):
+    DESTINATION = 0
+    ACK = 1

ltc/sim/process_output.py

@@ -5,7 +5,7 @@
 
 
 def no_transmission(args):
-    action, buffer_state, _, _, no_tx = args
+    action, buffer_state, _, _, no_tx, _ = args
     return jax.lax.cond(
         action == Actions.IDLE.value,
         lambda: jax.lax.cond(buffer_state == 0, idle_empty_buffer, idle_full_buffer, args),
@@ -21,34 +21,34 @@ def idle_empty_buffer(_):
 
 
 def idle_full_buffer(args):
-    _, _, ret_c, _, no_tx = args
+    _, _, ret_c, _, no_tx, _ = args
     reward = NO_TX_REWARD
     no_tx = no_tx + 1
     return reward, ret_c, no_tx
 
 
 def no_transmission_short(args):
-    _, buffer_state, ret_c, _, no_tx = args
+    _, buffer_state, ret_c, _, no_tx, _ = args
     reward = NO_TX_REWARD
     no_tx = jnp.where(buffer_state == 0, 0, no_tx + 1)
     return reward, ret_c, no_tx
 
 
 def no_transmission_long(args):
-    _, _, ret_c, _, no_tx = args
+    _, _, ret_c, _, no_tx, _ = args
     scale = jax.lax.min(1., (no_tx - SAFE_IDLE_PERIOD + 1) / PENALIZED_IDLE_PERIOD)
     reward = scale * NO_TX_PENALTY
     no_tx = no_tx + 1
     return reward, ret_c, no_tx
 
 
 def transmission(args):
-    _, _, _, channel_state, _ = args
+    _, _, _, channel_state, _, _ = args
     return jax.lax.cond(channel_state == 1, transmission_without_collision, transmission_with_collision, args)
 
 
 def transmission_with_collision(args):
-    _, _, ret_c, _, _ = args
+    _, _, ret_c, _, _, _ = args
     return jax.lax.cond(ret_c < MAX_RETRANSMISSION, retransmission, max_retransmission_collision, args)
 
 
@@ -60,15 +60,15 @@ def max_retransmission_collision(_):
 
 
 def retransmission(args):
-    _, _, ret_c, _, _ = args
+    _, _, ret_c, _, _, _ = args
     reward = COLLISION_PENALTY
     ret_c = ret_c + 1
     no_tx = 0
     return reward, ret_c, no_tx
 
 
 def transmission_without_collision(args):
-    _, buffer_state, _, _, _ = args
+    _, buffer_state, _, _, _, _ = args
     return jax.lax.cond(buffer_state > 0, successful_transmission, empty_buffer_transmission, args)
 
 
@@ -80,18 +80,61 @@ def empty_buffer_transmission(_):
 
 
 def successful_transmission(args):
-    _, _, ret_c, _, _ = args
+    _, _, ret_c, _, _, _ = args
     reward = TX_REWARD / (ret_c + 1)
     ret_c = 0
     no_tx = 0
     return reward, ret_c, no_tx
 
 
-def process_output_i(buffer_state, new_buffer_state, power_state, channel_state, obs, action, terminal):
-    _, _, ret_c, no_tx, _ = obs[-1]
-    args = (action, buffer_state, ret_c, channel_state, no_tx)
+def transmission_ack(args):
+    _, _, _, channel_state, _, _ = args
+    return jax.lax.cond(channel_state == 1, transmission_ack_without_collision, transmission_with_collision, args)
+
+
+def transmission_ack_without_collision(args):
+    _, _, _, _, _, ack_flag = args
+    return jax.lax.cond(ack_flag == 1, successful_ack_transmission, empty_ack_transmission, args)
+
+
+def empty_ack_transmission(args):
+    _, _, ret_c, _, no_tx, _ = args
+    reward = EMPTY_ACK_PENALTY
+    ret_c = ret_c
+    no_tx = no_tx + 1
+    return reward, ret_c, no_tx
+
+
+def successful_ack_transmission(args):
+    _, _, ret_c, _, _, _ = args
+    reward = ACK_REWARD / (ret_c + 1)
+    ret_c = 0
+    no_tx = 0
+    return reward, ret_c, no_tx
+
+
+def process_output_i(buffer_state, new_buffer_state, power_state, channel_state, obs, action, terminal,
+                     transmission_history):
+    _, _, ret_c, no_tx, _, my_id = obs[-1]
+    mask = (transmission_history[:, TransmissionIndex.DESTINATION.value] == my_id) & (
+            transmission_history[:, TransmissionIndex.ACK.value] == AckState.NOT_SENT.value)
+    matching_rows = jnp.nonzero(mask, size=mask.shape[0], fill_value=-1)[0]
+    print(matching_rows)
+    ack_flag = jnp.where(jnp.all(matching_rows == -1), -1, 1)
+    args = (action, buffer_state, ret_c, channel_state, no_tx, ack_flag)
+
+    reward, ret_c, no_tx = jax.lax.cond(
+        action == Actions.ACK.value,
+        transmission_ack,
+        lambda a: jax.lax.cond(
+            action == Actions.TX.value,
+            transmission,
+            no_transmission,
+            operand=a
+        ),
+        operand=args
+    )
 
-    reward, ret_c, no_tx = jax.lax.cond(action == Actions.TX.value, transmission, no_transmission, args)
     reward = jnp.where(terminal, 0., reward)
 
     channel_state = jnp.where(action == Actions.CS.value, channel_state, -1)
@@ -101,18 +144,24 @@ def process_output_i(buffer_state, new_buffer_state, power_state, channel_state,
             action == Actions.CS.value, power_state - CS_CONSUMPTION,
             jnp.where(
                 action == Actions.IDLE.value, power_state - IDLE_CONSUMPTION,
-                power_state
+                jnp.where(
+                    action == Actions.ACK.value, power_state - ACK_CONSUMPTION,
+                    power_state
+                )
             )
         )
     )
 
-    obs_t = jnp.array([new_buffer_state, channel_state, ret_c, no_tx, power])
+    obs_t = jnp.array([new_buffer_state, channel_state, ret_c, no_tx, power, my_id])
     obs = jnp.roll(obs, -1, axis=0)
     obs = obs.at[-1].set(obs_t)
 
     return obs, reward, power
 
 
-def process_output(buffer_states, new_buffer_states, power_states, channel_state, obs, actions, terminals):
+def process_output(buffer_states, new_buffer_states, power_states, channel_state, obs, actions, terminals, transmission_history):
     channel_states = jnp.full(buffer_states.shape[0], channel_state)
-    return jax.vmap(process_output_i)(buffer_states, new_buffer_states, power_states, channel_states, obs, actions, terminals)
+    return jax.vmap(process_output_i, in_axes=(0, 0, 0, 0, 0, 0, 0, None))(buffer_states, new_buffer_states,
+                                                                           power_states, channel_states, obs, actions,
+                                                                           terminals, transmission_history
+                                                                           )

ltc/sim/sim.py

@@ -1,6 +1,7 @@
 import jax.numpy as jnp
+import jax
 
-from ltc.sim.constants import Actions
+from ltc.sim.constants import *
 
 
 def channel_state_selector(actions):
@@ -11,10 +12,12 @@ def channel_state_selector(actions):
     int: -1 if more than one STA transmit, 1 if exactly one STA transmit, 0 if noone transmit at the moment.
     """
 
-    ones_count = jnp.sum(actions == Actions.TX.value)
+    transmitting = (actions == Actions.TX.value) | (actions == Actions.ACK.value)
+    active_count = jnp.sum(transmitting)
+
     return jnp.where(
-        ones_count > 1, -1,
-        jnp.where(ones_count == 1, 1, 0)
+        active_count > 1, -1,
+        jnp.where(active_count == 1, 1, 0)
     )
 
 
@@ -46,8 +49,52 @@ def add_new_frames(buffer_states, new_frames):
     return jnp.bitwise_or(buffer_states, new_frames)
 
 
-def simulate(buffer_states, new_frames, actions):
+def modify_transmission_history(actions, transmission_history, address):
+    args = (actions, transmission_history, address)
+    return jax.lax.cond(
+        jnp.any(actions == Actions.ACK.value),
+        ack_transmission,
+        lambda a: jax.lax.cond(
+            jnp.any(actions == Actions.TX.value),
+            data_transmission,
+            no_transmission,
+            operand=a
+        ),
+        operand=args
+    )
+
+
+def data_transmission(args):
+    actions, transmission_history, address = args
+    transmission_history_t = jnp.array([address, AckState.NOT_SENT.value])
+    transmission_history = jnp.roll(transmission_history, -1, axis=0)
+    transmission_history = transmission_history.at[-1].set(transmission_history_t)
+    return transmission_history
+
+
+def ack_transmission(args):
+    actions, transmission_history, address = args
+    sta_num = jnp.argmax(actions == Actions.ACK.value)
+    mask = (transmission_history[:, TransmissionIndex.DESTINATION.value] == sta_num) & (
+            transmission_history[:, TransmissionIndex.ACK.value] == AckState.NOT_SENT.value)
+    matching_rows = jnp.nonzero(mask, size=mask.shape[0], fill_value=-1)[0]
+    transmission_history = transmission_history.at[matching_rows[0], TransmissionIndex.ACK.value].set(
+        AckState.SENT.value)
+    return transmission_history
+
+
+def no_transmission(args):
+    actions, transmission_history, address = args
+    return transmission_history
+
+
+def simulate(key, buffer_states, new_frames, actions, transmission_history, nWifi):
     channel_state = channel_state_selector(actions)
+    address = jnp.where(channel_state == 1, jax.random.randint(key, shape=(), minval=0, maxval=nWifi), -1)
+    transmission_history = jnp.where(channel_state == 1,
+                                     modify_transmission_history(actions, transmission_history, address),
+                                     transmission_history)
     buffer_states = jnp.where(channel_state == 1, buffer_clearing(buffer_states, actions), buffer_states)
     buffer_states = add_new_frames(buffer_states, new_frames)
-    return buffer_states, channel_state
+
+    return buffer_states, channel_state, transmission_history

ltc/utils/scan_states.py

@@ -14,6 +14,7 @@ class Carry:
     traffic_states: ModelState
     buffer_states: jax.Array
     power_states: jax.Array
+    transmission_history: jax.Array
     channel_state: int
     key: jax.random.PRNGKey
     obs: jax.Array