<kaggle_start><code># # Deep Q-learner starter code
# **Work in progress!** Please forgive lack of clarity, bugs, and typos.
# This notebook aims to demonstrate creating a deep Q-learner using Keras and train it on the GFootball enviroment. It will focus on the coded need to train a deep q-learner agent, rather than theory. Hopefully it will provide a starting point developing models and testing them out on the GFootball enviroment, and show how to train and evaluate them using the Gym API. I have no idea at this point how well a Q-learner is likely to perform in this competition, though! It'll certainly take a lot of training and tweaking, and and a large replay buffer to perform well. See also [Convolutional Deep-Q learner](https://www.kaggle.com/garethjns/convolutional-deep-q-learner) for a convolutional version using the SMM environment wrapper.
# The code will be based on a package I've been working on while learning about RL myself. It's available here: [reinforcement-learning-keras](https://github.com/garethjns/reinforcement-learning-keras), and includes configs for solving other Gym enviroments including Cart-pole, Pong, Space invaders, Doom etc. It also handles running experiments with different agents, and some of the diffulties of working with Tensorflow. I've stripped a lot of that out here to try and keep the code simple. For those interested in theory, check the readme on github, there's a list of useful resources. There's an example using this package at the end of this notebook.
# GFootball environment.
# Some helper code
import collections
from typing import Union, Callable, List, Tuple, Iterable, Any
from dataclasses import dataclass
from IPython.display import Image
import matplotlib.pyplot as plt
import numpy as np
from tensorflow import keras
import tensorflow as tf
import seaborn as sns
import gym
import gfootball
import sys

sns.set()
# Training keras models in a loop with eager execution on causes memory leaks and terrible performance.
tf.compat.v1.disable_eager_execution()
sys.path.append("/kaggle/working/kaggle-football/")


# # Q-Learning
# Q-learners are off-policy temporal difference learners, have a look [here](https://lilianweng.github.io/lil-log/2018/02/19/a-long-peek-into-reinforcement-learning.html#model-transition-and-reward) for details and context within the RL world.
# ## Components
# - Agent
# - Model, usually a neural network
# - Replay buffer (the agent's memory)
# - Exploration method, such as episilon greedy
#
# ## Model
# We'll use the GFootball enviroment with the Simple115StateWrapper which returns observations (state) as a (115,) array. This will be used as the input to a very simple neural network wtih 3 layers, which will output an estimated value for every possible action. The agent will either select the action with the highest value as estimated from this model, or randomly (chance depending on epsilon) select a random action.
# To make learning more stable with deep-q learners, two copies of this model are used by the agent.
# - "Action" model This copy is updated every training step and is used to estimate the value of actions for the current state.
# - "Target" model" The copy is updated by copying the weights of the value model, this is done less frequently in order to keep the training target more stable.
#
# The model is defined in this class, feel free to play with the arcitecture! Also note that unit_scale is set to make the model very small, this makes running in the notebook easier, but should be increased for real training.
class DenseNN:
    def __init__(
        self,
        observation_shape: List[int],
        n_actions: int,
        output_activation: Union[None, str] = None,
        unit_scale: int = 1,
        learning_rate: float = 0.0001,
        opt: str = "Adam",
    ) -> None:
        """
        :param observation_shape: Tuple specifying input shape.
        :param n_actions: Int specifying number of outputs
        :param output_activation: Activation function for output. Eg.
                                  None for value estimation (off-policy methods).
        :param unit_scale: Multiplier for all units in FC layers in network
                           (not used here at the moment).
        :param opt: Keras optimiser to use. Should be string.
                    This is to avoid storing TF/Keras objects here.
        :param learning_rate: Learning rate for optimiser.
        """
        self.observation_shape = observation_shape
        self.n_actions = n_actions
        self.unit_scale = unit_scale
        self.output_activation = output_activation
        self.learning_rate = learning_rate
        self.opt = opt

    def _model_architecture(self) -> Tuple[keras.layers.Layer, keras.layers.Layer]:
        frame_input = keras.layers.Input(name="input", shape=self.observation_shape)
        fc1 = keras.layers.Dense(
            int(self.observation_shape[0] / 1.5 * self.unit_scale),
            name="fc1",
            activation="relu",
        )(frame_input)
        fc2 = keras.layers.Dense(
            int(self.observation_shape[0] / 3 * self.unit_scale),
            name="fc2",
            activation="relu",
        )(fc1)
        fc3 = keras.layers.Dense(self.n_actions * 2, name="fc3", activation="relu")(fc2)
        action_output = keras.layers.Dense(
            units=self.n_actions, name="output", activation=self.output_activation
        )(fc3)
        return frame_input, action_output

    def compile(
        self, model_name: str = "model", loss: Union[str, Callable] = "mse"
    ) -> keras.Model:
        """
        Compile a copy of the model using the provided loss.
        :param model_name: Name of model
        :param loss: Model loss. Default 'mse'. Can be custom callable.
        """
        # Get optimiser
        if self.opt.lower() == "adam":
            opt = keras.optimizers.Adam
        elif self.opt.lower() == "rmsprop":
            opt = keras.optimizers.RMSprop
        else:
            raise ValueError(f"Invalid optimiser {self.opt}")
        state_input, action_output = self._model_architecture()
        model = keras.Model(
            inputs=[state_input], outputs=[action_output], name=model_name
        )
        model.compile(optimizer=opt(learning_rate=self.learning_rate), loss=loss)
        return model

    def plot(self, model_name: str = "model") -> None:
        keras.utils.plot_model(
            self.compile(model_name), to_file=f"{model_name}.png", show_shapes=True
        )
        plt.show()


DenseNN(observation_shape=(115,), unit_scale=0.25, n_actions=19).plot()
Image(filename="model.png")


# ## Replay buffer
# The replay buffer records previous state, actions, and rewards. When the agent is updated, it samples a batch from the replay buffer to create the training set. These samples can come from different episodes if the buffer is long enough. This is necessary to reduce the correlation in the training data, which would be extremely high if, for example, the model was trained with just the last observation on each step. Neural networks don't do well with highly correlated training data.
@dataclass
class ContinuousBuffer:
    buffer_size: int = 250

    def __post_init__(self) -> None:
        self._state_queue = collections.deque(maxlen=self.buffer_size)
        self._other_queue = collections.deque(maxlen=self.buffer_size)
        self.queue = collections.deque(maxlen=self.buffer_size)

    def __len__(self) -> int:
        return self.n if (self.n > 0) else 0

    @property
    def full(self) -> bool:
        return len(self._state_queue) == self.buffer_size

    @property
    def n(self) -> int:
        return len(self._state_queue) - 1

    def append(self, items: Tuple[Any, int, float, bool]) -> None:
        """
        :param items: Tuple containing (s, a, r, d).
        """
        self._state_queue.append(items[0])
        self._other_queue.append(items[1::])

    def get_batch(
        self, idxs: Iterable[int]
    ) -> Tuple[
        List[np.ndarray], List[np.ndarray], List[float], List[bool], List[np.ndarray]
    ]:
        ss = [self._state_queue[i] for i in idxs]
        ss_ = [self._state_queue[i + 1] for i in idxs]
        ard = [self._other_queue[i] for i in idxs]
        aa = [a for (a, _, _) in ard]
        rr = [r for (_, r, _) in ard]
        dd = [d for (_, _, d) in ard]
        return ss, aa, rr, dd, ss_

    def sample_batch(
        self, n: int
    ) -> Tuple[
        List[np.ndarray], List[np.ndarray], List[float], List[bool], List[np.ndarray]
    ]:
        if n > self.n:
            raise ValueError
        idxs = np.random.randint(0, self.n, n)
        return self.get_batch(idxs)


# ## Epsilon greedy
# Epsilon greedy is used to force the agent to explore the environment. When the agent is asked for an action, it draws a random sample between 0-1. If this value is >epsilon, the agent simply selects a random action. Typically epsilon is set high at the beginning of training, and is slowly decayed over time. This class tracks it and handles the random action selection.
@dataclass
class EpsilonGreedy:
    """
    Handles epsilon-greedy action selection, decay of epsilon during training.
    :param eps_initial: Initial epsilon value.
    :param decay: Decay rate in percent (should be positive to decay).
    :param decay_schedule: 'linear' or 'compound'.
    :param eps_min: The min value epsilon can fall to.
    :param state: Random state, used to pick between the greedy or random options.
    """

    eps_initial: float = 0.2
    decay: float = 0.0001
    decay_schedule: str = "compound"
    eps_min: float = 0.01
    state = None

    def __post_init__(self) -> None:
        self._step: int = 0
        self.eps_current = self.eps_initial
        valid_decay = ("linear", "compound")
        if self.decay_schedule.lower() not in valid_decay:
            raise ValueError(
                f"Invalid decay schedule {self.decay_schedule}. "
                "Pick from {valid_decay}."
            )
        self._set_random_state()

    def _set_random_state(self) -> None:
        self._state = np.random.RandomState(self.state)

    def _linear_decay(self) -> float:
        return self.eps_current - self.decay

    def _compound_decay(self) -> float:
        return self.eps_current - self.eps_current * self.decay

    def _decay(self):
        new_eps = np.nan
        if self.decay_schedule.lower() == "linear":
            new_eps = self._linear_decay()
        if self.decay_schedule.lower() == "compound":
            new_eps = self._compound_decay()
        self._step += 1
        return max(self.eps_min, new_eps)

    def select(
        self, greedy_option: Callable, random_option: Callable, training: bool = False
    ) -> Any:
        """
        Apply epsilon greedy selection.
        If training, decay epsilon, and return selected option.
        If not training, just return greedy_option.
        Use of lambdas is to avoid unnecessarily picking between
        two pre-computed options.
        :param greedy_option: Function to evaluate if random option
                              is NOT picked.
        :param random_option: Function to evaluate if random option
                              IS picked.
        :param training: Bool indicating if call is during training
                         and to use epsilon greedy and decay.
        :return: Evaluated selected option.
        """
        if training:
            self.eps_current = self._decay()
            if self._state.random() < self.eps_current:
                return random_option()
        return greedy_option()


# ## Agent
# The agent estimates actions values from states, and learns from its experience. Most of the magic happens in the .update_model, where the problem is formulated into something that can be learned by the DenseNN model.
@dataclass
class DeepQAgent:
    replay_buffer: ContinuousBuffer
    eps: EpsilonGreedy
    model_architecture: DenseNN
    name: str = "DQNAgent"
    double: bool = False
    noisy: bool = False
    gamma: float = 0.99
    replay_buffer_samples: int = 75
    final_reward: Union[float, None] = None

    def __post_init__(self) -> None:
        self._build_model()

    def _build_model(self) -> None:
        """
        Prepare two of the same model.
        The action model is used to pick actions and the target model
        is used to predict value of Q(s', a). Action model
        weights are updated on every buffer sample + training step.
        The target model is never directly trained, but it's
        weights are updated to match the action model at the end of
        each episode.
        """
        self._action_model = self.model_architecture.compile(
            model_name="action_model", loss="mse"
        )
        self._target_model = self.model_architecture.compile(
            model_name="target_model", loss="mse"
        )

    def transform(self, s: np.ndarray) -> np.ndarray:
        """Check input shape, add Row dimension if required."""
        if len(s.shape) < len(self._action_model.input.shape):
            s = np.expand_dims(s, 0)
        return s

    def update_experience(self, s: np.ndarray, a: int, r: float, d: bool) -> None:
        """
        First the most recent step is added to the buffer.
        Note that s' isn't saved because there's no need.
        It'll be added next step. s' for any s is always index + 1 in
        the buffer.
        """
        # Add s, a, r, d to experience buffer
        self.replay_buffer.append((s, a, r, d))

    def update_model(self) -> None:
        """
        Sample a batch from the replay buffer, calculate targets using
        target model, and train action model.
        If the buffer is below its minimum size, no training is done.
        If the buffer has reached its minimum size, a training batch
        from the replay buffer and the action model is updated.
        This update samples random (s, a, r, s') sets from the buffer
        and calculates the discounted reward for each set.
        The value of the actions at states s and s' are predicted from
        the value model. The action model is updated using these value
        predictions as the targets. The value of performed action is
        updated with the discounted reward (using its value prediction
        at s'). ie. x=s, y=[action value 1, action value 2].
        """
        # If buffer isn't full, don't train
        if not self.replay_buffer.full:
            return
        # Else sample batch from buffer
        ss, aa, rr, dd, ss_ = self.replay_buffer.sample_batch(
            self.replay_buffer_samples
        )
        # Calculate estimated S,A values for current states and next states.
        # These are stacked together first to avoid making two separate
        # predict calls (which is slow on GPU).
        ss = np.array(ss)
        ss_ = np.array(ss_)
        y_now_and_future = self._target_model.predict_on_batch(np.vstack((ss, ss_)))
        # Separate again
        y_now = y_now_and_future[0 : self.replay_buffer_samples]
        y_future = y_now_and_future[self.replay_buffer_samples : :]
        # Update rewards where not done with y_future predictions
        dd_mask = np.array(dd, dtype=bool).squeeze()
        rr = np.array(rr, dtype=float).squeeze()
        # Gather max action indexes and update relevant actions in y
        if self.double:
            # If using double dqn select best actions using the action model,
            # but the value of those action using the
            # target model (already have in y_future).
            y_future_action_model = self._action_model.predict_on_batch(ss_)
            selected_actions = np.argmax(y_future_action_model[~dd_mask, :], axis=1)
        else:
            # If normal dqn select targets using target model,
            # and value of those from target model too
            selected_actions = np.argmax(y_future[~dd_mask, :], axis=1)
        # Update reward values with estimated values (where not done)
        # and final rewards (where done)
        rr[~dd_mask] += y_future[~dd_mask, selected_actions]
        if self.final_reward is not None:
            # If self.final_reward is set, set done cases to this value.
            # Else leave as observed reward.
            rr[dd_mask] = self.final_reward
        aa = np.array(aa, dtype=int)
        np.put_along_axis(y_now, aa.reshape(-1, 1), rr.reshape(-1, 1), axis=1)
        # Fit model with updated y_now values
        self._action_model.train_on_batch(ss, y_now)

    def get_best_action(self, s: np.ndarray) -> np.ndarray:
        """
        Get best action(s) from model - the one with the highest predicted value.
        :param s: A single or multiple rows of state observations.
        :return: The selected action.
        """
        preds = self._action_model.predict(self.transform(s))
        return np.argmax(preds)

    def get_action(self, s: np.ndarray, training: bool = False) -> int:
        """
        Get an action using get_best_action or epsilon greedy.
        Epsilon decays every time a random action is chosen.
        :param s: The raw state observation.
        :param training: Bool to indicate whether or not to use this
                         experience to update the model. If False, just
                         returns best action.
        :return: The selected action.
        """
        action = self.eps.select(
            greedy_option=lambda: self.get_best_action(s),
            random_option=lambda: self.env.action_space.sample(),
            training=training,
        )
        return action

    def update_target_model(self) -> None:
        """
        Update the value model with the weights of the action model
        (which is updated each step).
        The value model is updated less often to aid stability.
        """
        self._target_model.set_weights(self._action_model.get_weights())

    def after_episode_update(self) -> None:
        """Value model synced with action model at the end of each episode."""
        self.update_target_model()

    def _discounted_reward(
        self, reward: float, estimated_future_action_rewards: np.ndarray
    ) -> float:
        """
        Use this to define the discounted reward for unfinished episodes,
        default is 1 step TD.
        """
        return reward + self.gamma * np.max(estimated_future_action_rewards)

    def _get_reward(
        self, reward: float, estimated_future_action_rewards: np.ndarray, done: bool
    ) -> float:
        """
        Calculate discounted reward for a single step.
        :param reward: Last real reward.
        :param estimated_future_action_rewards: Estimated future values
                                                of actions taken on next step.
        :param done: Flag indicating if this is the last step on an episode.
        :return: Reward.
        """
        if done:
            # If done, reward is just this step. Can finish because agent has won or lost.
            return self._final_reward(reward)
        else:
            # Otherwise, it's the reward plus the predicted max value of next action
            return self._discounted_reward(reward, estimated_future_action_rewards)


# # Play single episode
# An episode refers to one play through an enviroment. It will take many episodes to train an agent
def play_episode(
    env: gym.Env, agent, n_steps: int = 10, pr: bool = False, training: bool = False
) -> Tuple[List[float], List[int]]:
    episode_rewards = []
    episode_actions = []
    obs = env.reset()
    done = False
    for s in range(n_steps):
        if done:
            break
        # Select action
        action = agent.get_action(obs)
        episode_actions.append(action)
        # Take action
        prev_obs = obs
        obs, reward, done, info = env.step(action)
        episode_rewards.append(reward)
        # Update model
        if training:
            agent.update_experience(s=prev_obs, a=action, r=reward, d=done)
            agent.update_model()
        if pr:
            print(f"Step {s}: Action taken {action}, " f"reward recieved {reward}")
        last_step = s
    if training:
        agent.after_episode_update()
    return episode_rewards, episode_actions


env = gym.make("GFootball-11_vs_11_kaggle-simple115v2-v0")
agent = DeepQAgent(
    # Realistically this should be more like 100,000!
    replay_buffer=ContinuousBuffer(buffer_size=2000),
    eps=EpsilonGreedy(),
    model_architecture=DenseNN(
        observation_shape=env.observation_space.shape,
        # This is also set too low to make running in a notebook easier
        unit_scale=0.3,
        n_actions=env.action_space.n,
    ),
)
_ = play_episode(env, agent, n_steps=100, pr=True, training=True)


#
# # Play multiple episodes
def run_multiple_episodes(
    env: gym.Env,
    agent,
    n_episodes: int = 10,
    n_steps: int = 10,
    pr: bool = False,
    training: bool = False,
) -> List[float]:
    total_episode_rewards = []
    for ep in range(n_episodes):
        episode_rewards, _ = play_episode(env, agent, n_steps, pr=False, training=True)
        total_episode_rewards.append(sum(episode_rewards))
        if pr:
            print(
                f"Episode {ep} finished after {len(episode_rewards)} "
                f"steps, total reward: {sum(episode_rewards)}"
            )
    return total_episode_rewards


env = gym.make("GFootball-11_vs_11_kaggle-simple115v2-v0")
agent = DeepQAgent(
    replay_buffer=ContinuousBuffer(buffer_size=1000),
    eps=EpsilonGreedy(eps_initial=0.7),
    model_architecture=DenseNN(
        observation_shape=env.observation_space.shape,
        unit_scale=0.25,
        n_actions=env.action_space.n,
    ),
)
# Both n_episodes and n_steps need to
reward_history = run_multiple_episodes(
    env,
    agent,
    n_episodes=10,
    # Use 3000 for 90 minute games
    n_steps=3000,
    training=True,
    pr=True,
)
plt.plot(reward_history)
plt.xlabel("Episode")
plt.ylabel("Reward")
plt.show()


# # Compare to a random agent
class RandomModel:
    def __init__(self, n_actions: int):
        self.n_actions = n_actions

    def predict(self) -> int:
        return int(np.random.randint(0, self.n_actions, 1))


@dataclass
class RandomAgent:
    """A random agent that acts randomly and learns nothing."""

    n_actions: int
    name: str = "RandomAgent"

    def __post_init__(self) -> None:
        self._build_model()

    def _build_model(self) -> None:
        """Set model function."""
        self.model = RandomModel(self.n_actions)

    def update_model(self, *args, **kwargs) -> None:
        """No model to update."""
        pass

    def get_action(self, s: Any, **kwargs) -> int:
        return self.model.predict()


env = gym.make("GFootball-11_vs_11_kaggle-simple115v2-v0")
obs = env.reset()
random_agent = RandomAgent(n_actions=env.action_space.n)
ra_reward_history = run_multiple_episodes(
    env, agent, n_episodes=25, n_steps=3000, training=False
)
dqn_reward_history = run_multiple_episodes(
    env, agent, n_episodes=25, n_steps=3000, training=False
)
plt.plot(ra_reward_history, label="Random")
plt.plot(dqn_reward_history, label="DQN")
plt.xlabel("Episode")
plt.ylabel("Reward")
plt.legend(title="Agent")
plt.show()


# # Prepare submission
# This is covered in a bit more detail in the next notebook: https://www.kaggle.com/garethjns/convolutional-deep-q-learner . In short, we need to create a single file containing the agent components we need for inference, including the serialsed model weights. We also need to add to agent function that acts given an environment observation.
# ## Get model weights
import zlib
import pickle

weights_bytes = zlib.compress(pickle.dumps(agent._action_model.get_weights()))

# ## Create submission.py
# Copy and paste the contents of weight_bytes into WEIGHTS_BYTES
import zlib
import pickle
import tensorflow as tf
import numpy as np
from tensorflow import keras
from typing import List, Tuple, Union, Callable
from gfootball.env.wrappers import Simple115StateWrapper


class DenseNN:
    def __init__(
        self,
        observation_shape: List[int],
        n_actions: int,
        output_activation: Union[None, str] = None,
        unit_scale: int = 1,
        learning_rate: float = 0.0001,
        opt: str = "Adam",
    ) -> None:
        """
        :param observation_shape: Tuple specifying input shape.
        :param n_actions: Int specifying number of outputs
        :param output_activation: Activation function for output. Eg.
                                  None for value estimation (off-policy methods).
        :param unit_scale: Multiplier for all units in FC layers in network
                           (not used here at the moment).
        :param opt: Keras optimiser to use. Should be string.
                    This is to avoid storing TF/Keras objects here.
        :param learning_rate: Learning rate for optimiser.
        """
        self.observation_shape = observation_shape
        self.n_actions = n_actions
        self.unit_scale = unit_scale
        self.output_activation = output_activation
        self.learning_rate = learning_rate
        self.opt = opt

    def _model_architecture(self) -> Tuple[keras.layers.Layer, keras.layers.Layer]:
        frame_input = keras.layers.Input(name="input", shape=self.observation_shape)
        fc1 = keras.layers.Dense(
            int(self.observation_shape[0] / 1.5 * self.unit_scale),
            name="fc1",
            activation="relu",
        )(frame_input)
        fc2 = keras.layers.Dense(
            int(self.observation_shape[0] / 3 * self.unit_scale),
            name="fc2",
            activation="relu",
        )(fc1)
        fc3 = keras.layers.Dense(self.n_actions * 2, name="fc3", activation="relu")(fc2)
        action_output = keras.layers.Dense(
            units=self.n_actions, name="output", activation=self.output_activation
        )(fc3)
        return frame_input, action_output

    def compile(
        self, model_name: str = "model", loss: Union[str, Callable] = "mse"
    ) -> keras.Model:
        """
        Compile a copy of the model using the provided loss.
        :param model_name: Name of model
        :param loss: Model loss. Default 'mse'. Can be custom callable.
        """
        # Get optimiser
        if self.opt.lower() == "adam":
            opt = keras.optimizers.Adam
        elif self.opt.lower() == "rmsprop":
            opt = keras.optimizers.RMSprop
        else:
            raise ValueError(f"Invalid optimiser {self.opt}")
        state_input, action_output = self._model_architecture()
        model = keras.Model(
            inputs=[state_input], outputs=[action_output], name=model_name
        )
        model.compile(optimizer=opt(learning_rate=self.learning_rate), loss=loss)
        return model

    def plot(self, model_name: str = "model") -> None:
        keras.utils.plot_model(
            self.compile(model_name), to_file=f"{model_name}.png", show_shapes=True
        )
        plt.show()


tf_mod = DenseNN(
    observation_shape=(115,),
    # This needs to match the number from training
    unit_scale=0.25,
    n_actions=19,
).compile()
tf_mod.set_weights(pickle.loads(zlib.decompress(WEIGHTS_BYTES)))


def agent(obs):
    # Use the existing model and obs buffer on each call to agent
    global tf_mod
    # Convert these to the same output as the Simple115StateWrapper we used in training
    obs = Simple115StateWrapper.convert_observation(
        obs["players_raw"], fixed_positions=False
    )
    # Predict actions from keras model
    actions = tf_mod.predict(obs)
    action = np.argmax(actions)
    return [action]


# # Test submission.py
# The written submission file can be tested with the following block.
# See here: https://github.com/garethjns/kaggle-football/blob/main/scripts/debug_agent.py for a version that can be used to run locally, but also maintain debugability.
from kaggle_environments import make

env = make(
    "football", configuration={"save_video": True, "scenario_name": "11_vs_11_kaggle"}
)
# Define players
left_player = "submission.py"  # A custom agent, eg. random_agent.py or example_agent.py
right_player = "run_right"  # eg. A built in 'AI' agent or the agent again
output = env.run([left_player, right_player])
print(
    f"Final score: {sum([r['reward'] for r in output[0]])} : {sum([r['reward'] for r in output[1]])}"
)
env.render(mode="human", width=800, height=600)

# # Conclusions and next steps
# The DQN performs currently extremely poorly, which is expected given the totally inadequate training done so far here. However, it seems to be running, so that's a start.
# A few things to try:
# - Increase number of training iterations
# - Increase size of replay buffer
# - Increase number of units in network (set unit_scale higher)
# - Increase depth of NN (modify architecture)
# # Appendix
# ## Training with reinforcement-learning-keras
# Follows the same logic as above, but is a bit more fancy. https://github.com/garethjns/reinforcement-learning-keras
import gfootball
from reinforcement_learning_keras.agents.components.history.training_history import (
    TrainingHistory,
)
from reinforcement_learning_keras.agents.components.replay_buffers.continuous_buffer import (
    ContinuousBuffer,
)
from reinforcement_learning_keras.agents.models.dense_nn import DenseNN
from reinforcement_learning_keras.agents.q_learning.deep_q_agent import DeepQAgent
from reinforcement_learning_keras.agents.q_learning.exploration.epsilon_greedy import (
    EpsilonGreedy,
)

agent = DeepQAgent(
    name="deep_q",
    model_architecture=DenseNN(observation_shape=(115,), n_actions=19),
    replay_buffer=ContinuousBuffer(buffer_size=30000),
    env_spec="GFootball-11_vs_11_kaggle-simple115v2-v0",
    eps=EpsilonGreedy(
        eps_initial=0.9, decay=0.0001, eps_min=0.01, decay_schedule="linear"
    ),
    training_history=TrainingHistory(
        agent_name="deep_q", rolling_average=1, plotting_on=True, plot_every=5
    ),
)
agent.train(
    verbose=True,
    render=False,
    n_episodes=10,
    max_episode_steps=3000,
    update_every=2,
    checkpoint_every=1,
)
<|endoftext|>
[{"Id": 45605603, "ScriptId": 12495311, "ParentScriptVersionId": NaN, "ScriptLanguageId": 9, "AuthorUserId": 4999684, "CreationDate": "10/27/2020 12:44:35", "VersionNumber": 3.0, "Title": "Wstawic115v3", "EvaluationDate": "10/27/2020", "IsChange": true, "TotalLines": 794.0, "LinesInsertedFromPrevious": 4.0, "LinesChangedFromPrevious": 0.0, "LinesUnchangedFromPrevious": 790.0, "LinesInsertedFromFork": 4.0, "LinesDeletedFromFork": 4.0, "LinesChangedFromFork": 0.0, "LinesUnchangedFromFork": 790.0, "TotalVotes": 0}]
# # Deep Q-learner starter code
# **Work in progress!** Please forgive lack of clarity, bugs, and typos.
# This notebook aims to demonstrate creating a deep Q-learner using Keras and train it on the GFootball enviroment. It will focus on the coded need to train a deep q-learner agent, rather than theory. Hopefully it will provide a starting point developing models and testing them out on the GFootball enviroment, and show how to train and evaluate them using the Gym API. I have no idea at this point how well a Q-learner is likely to perform in this competition, though! It'll certainly take a lot of training and tweaking, and and a large replay buffer to perform well. See also [Convolutional Deep-Q learner](https://www.kaggle.com/garethjns/convolutional-deep-q-learner) for a convolutional version using the SMM environment wrapper.
# The code will be based on a package I've been working on while learning about RL myself. It's available here: [reinforcement-learning-keras](https://github.com/garethjns/reinforcement-learning-keras), and includes configs for solving other Gym enviroments including Cart-pole, Pong, Space invaders, Doom etc. It also handles running experiments with different agents, and some of the diffulties of working with Tensorflow. I've stripped a lot of that out here to try and keep the code simple. For those interested in theory, check the readme on github, there's a list of useful resources. There's an example using this package at the end of this notebook.
# GFootball environment.
# Some helper code
import collections
from typing import Union, Callable, List, Tuple, Iterable, Any
from dataclasses import dataclass
from IPython.display import Image
import matplotlib.pyplot as plt
import numpy as np
from tensorflow import keras
import tensorflow as tf
import seaborn as sns
import gym
import gfootball
import sys

sns.set()
# Training keras models in a loop with eager execution on causes memory leaks and terrible performance.
tf.compat.v1.disable_eager_execution()
sys.path.append("/kaggle/working/kaggle-football/")


# # Q-Learning
# Q-learners are off-policy temporal difference learners, have a look [here](https://lilianweng.github.io/lil-log/2018/02/19/a-long-peek-into-reinforcement-learning.html#model-transition-and-reward) for details and context within the RL world.
# ## Components
# - Agent
# - Model, usually a neural network
# - Replay buffer (the agent's memory)
# - Exploration method, such as episilon greedy
#
# ## Model
# We'll use the GFootball enviroment with the Simple115StateWrapper which returns observations (state) as a (115,) array. This will be used as the input to a very simple neural network wtih 3 layers, which will output an estimated value for every possible action. The agent will either select the action with the highest value as estimated from this model, or randomly (chance depending on epsilon) select a random action.
# To make learning more stable with deep-q learners, two copies of this model are used by the agent.
# - "Action" model This copy is updated every training step and is used to estimate the value of actions for the current state.
# - "Target" model" The copy is updated by copying the weights of the value model, this is done less frequently in order to keep the training target more stable.
#
# The model is defined in this class, feel free to play with the arcitecture! Also note that unit_scale is set to make the model very small, this makes running in the notebook easier, but should be increased for real training.
class DenseNN:
    def __init__(
        self,
        observation_shape: List[int],
        n_actions: int,
        output_activation: Union[None, str] = None,
        unit_scale: int = 1,
        learning_rate: float = 0.0001,
        opt: str = "Adam",
    ) -> None:
        """
        :param observation_shape: Tuple specifying input shape.
        :param n_actions: Int specifying number of outputs
        :param output_activation: Activation function for output. Eg.
                                  None for value estimation (off-policy methods).
        :param unit_scale: Multiplier for all units in FC layers in network
                           (not used here at the moment).
        :param opt: Keras optimiser to use. Should be string.
                    This is to avoid storing TF/Keras objects here.
        :param learning_rate: Learning rate for optimiser.
        """
        self.observation_shape = observation_shape
        self.n_actions = n_actions
        self.unit_scale = unit_scale
        self.output_activation = output_activation
        self.learning_rate = learning_rate
        self.opt = opt

    def _model_architecture(self) -> Tuple[keras.layers.Layer, keras.layers.Layer]:
        frame_input = keras.layers.Input(name="input", shape=self.observation_shape)
        fc1 = keras.layers.Dense(
            int(self.observation_shape[0] / 1.5 * self.unit_scale),
            name="fc1",
            activation="relu",
        )(frame_input)
        fc2 = keras.layers.Dense(
            int(self.observation_shape[0] / 3 * self.unit_scale),
            name="fc2",
            activation="relu",
        )(fc1)
        fc3 = keras.layers.Dense(self.n_actions * 2, name="fc3", activation="relu")(fc2)
        action_output = keras.layers.Dense(
            units=self.n_actions, name="output", activation=self.output_activation
        )(fc3)
        return frame_input, action_output

    def compile(
        self, model_name: str = "model", loss: Union[str, Callable] = "mse"
    ) -> keras.Model:
        """
        Compile a copy of the model using the provided loss.
        :param model_name: Name of model
        :param loss: Model loss. Default 'mse'. Can be custom callable.
        """
        # Get optimiser
        if self.opt.lower() == "adam":
            opt = keras.optimizers.Adam
        elif self.opt.lower() == "rmsprop":
            opt = keras.optimizers.RMSprop
        else:
            raise ValueError(f"Invalid optimiser {self.opt}")
        state_input, action_output = self._model_architecture()
        model = keras.Model(
            inputs=[state_input], outputs=[action_output], name=model_name
        )
        model.compile(optimizer=opt(learning_rate=self.learning_rate), loss=loss)
        return model

    def plot(self, model_name: str = "model") -> None:
        keras.utils.plot_model(
            self.compile(model_name), to_file=f"{model_name}.png", show_shapes=True
        )
        plt.show()


DenseNN(observation_shape=(115,), unit_scale=0.25, n_actions=19).plot()
Image(filename="model.png")


# ## Replay buffer
# The replay buffer records previous state, actions, and rewards. When the agent is updated, it samples a batch from the replay buffer to create the training set. These samples can come from different episodes if the buffer is long enough. This is necessary to reduce the correlation in the training data, which would be extremely high if, for example, the model was trained with just the last observation on each step. Neural networks don't do well with highly correlated training data.
@dataclass
class ContinuousBuffer:
    buffer_size: int = 250

    def __post_init__(self) -> None:
        self._state_queue = collections.deque(maxlen=self.buffer_size)
        self._other_queue = collections.deque(maxlen=self.buffer_size)
        self.queue = collections.deque(maxlen=self.buffer_size)

    def __len__(self) -> int:
        return self.n if (self.n > 0) else 0

    @property
    def full(self) -> bool:
        return len(self._state_queue) == self.buffer_size

    @property
    def n(self) -> int:
        return len(self._state_queue) - 1

    def append(self, items: Tuple[Any, int, float, bool]) -> None:
        """
        :param items: Tuple containing (s, a, r, d).
        """
        self._state_queue.append(items[0])
        self._other_queue.append(items[1::])

    def get_batch(
        self, idxs: Iterable[int]
    ) -> Tuple[
        List[np.ndarray], List[np.ndarray], List[float], List[bool], List[np.ndarray]
    ]:
        ss = [self._state_queue[i] for i in idxs]
        ss_ = [self._state_queue[i + 1] for i in idxs]
        ard = [self._other_queue[i] for i in idxs]
        aa = [a for (a, _, _) in ard]
        rr = [r for (_, r, _) in ard]
        dd = [d for (_, _, d) in ard]
        return ss, aa, rr, dd, ss_

    def sample_batch(
        self, n: int
    ) -> Tuple[
        List[np.ndarray], List[np.ndarray], List[float], List[bool], List[np.ndarray]
    ]:
        if n > self.n:
            raise ValueError
        idxs = np.random.randint(0, self.n, n)
        return self.get_batch(idxs)


# ## Epsilon greedy
# Epsilon greedy is used to force the agent to explore the environment. When the agent is asked for an action, it draws a random sample between 0-1. If this value is >epsilon, the agent simply selects a random action. Typically epsilon is set high at the beginning of training, and is slowly decayed over time. This class tracks it and handles the random action selection.
@dataclass
class EpsilonGreedy:
    """
    Handles epsilon-greedy action selection, decay of epsilon during training.
    :param eps_initial: Initial epsilon value.
    :param decay: Decay rate in percent (should be positive to decay).
    :param decay_schedule: 'linear' or 'compound'.
    :param eps_min: The min value epsilon can fall to.
    :param state: Random state, used to pick between the greedy or random options.
    """

    eps_initial: float = 0.2
    decay: float = 0.0001
    decay_schedule: str = "compound"
    eps_min: float = 0.01
    state = None

    def __post_init__(self) -> None:
        self._step: int = 0
        self.eps_current = self.eps_initial
        valid_decay = ("linear", "compound")
        if self.decay_schedule.lower() not in valid_decay:
            raise ValueError(
                f"Invalid decay schedule {self.decay_schedule}. "
                "Pick from {valid_decay}."
            )
        self._set_random_state()

    def _set_random_state(self) -> None:
        self._state = np.random.RandomState(self.state)

    def _linear_decay(self) -> float:
        return self.eps_current - self.decay

    def _compound_decay(self) -> float:
        return self.eps_current - self.eps_current * self.decay

    def _decay(self):
        new_eps = np.nan
        if self.decay_schedule.lower() == "linear":
            new_eps = self._linear_decay()
        if self.decay_schedule.lower() == "compound":
            new_eps = self._compound_decay()
        self._step += 1
        return max(self.eps_min, new_eps)

    def select(
        self, greedy_option: Callable, random_option: Callable, training: bool = False
    ) -> Any:
        """
        Apply epsilon greedy selection.
        If training, decay epsilon, and return selected option.
        If not training, just return greedy_option.
        Use of lambdas is to avoid unnecessarily picking between
        two pre-computed options.
        :param greedy_option: Function to evaluate if random option
                              is NOT picked.
        :param random_option: Function to evaluate if random option
                              IS picked.
        :param training: Bool indicating if call is during training
                         and to use epsilon greedy and decay.
        :return: Evaluated selected option.
        """
        if training:
            self.eps_current = self._decay()
            if self._state.random() < self.eps_current:
                return random_option()
        return greedy_option()


# ## Agent
# The agent estimates actions values from states, and learns from its experience. Most of the magic happens in the .update_model, where the problem is formulated into something that can be learned by the DenseNN model.
@dataclass
class DeepQAgent:
    replay_buffer: ContinuousBuffer
    eps: EpsilonGreedy
    model_architecture: DenseNN
    name: str = "DQNAgent"
    double: bool = False
    noisy: bool = False
    gamma: float = 0.99
    replay_buffer_samples: int = 75
    final_reward: Union[float, None] = None

    def __post_init__(self) -> None:
        self._build_model()

    def _build_model(self) -> None:
        """
        Prepare two of the same model.
        The action model is used to pick actions and the target model
        is used to predict value of Q(s', a). Action model
        weights are updated on every buffer sample + training step.
        The target model is never directly trained, but it's
        weights are updated to match the action model at the end of
        each episode.
        """
        self._action_model = self.model_architecture.compile(
            model_name="action_model", loss="mse"
        )
        self._target_model = self.model_architecture.compile(
            model_name="target_model", loss="mse"
        )

    def transform(self, s: np.ndarray) -> np.ndarray:
        """Check input shape, add Row dimension if required."""
        if len(s.shape) < len(self._action_model.input.shape):
            s = np.expand_dims(s, 0)
        return s

    def update_experience(self, s: np.ndarray, a: int, r: float, d: bool) -> None:
        """
        First the most recent step is added to the buffer.
        Note that s' isn't saved because there's no need.
        It'll be added next step. s' for any s is always index + 1 in
        the buffer.
        """
        # Add s, a, r, d to experience buffer
        self.replay_buffer.append((s, a, r, d))

    def update_model(self) -> None:
        """
        Sample a batch from the replay buffer, calculate targets using
        target model, and train action model.
        If the buffer is below its minimum size, no training is done.
        If the buffer has reached its minimum size, a training batch
        from the replay buffer and the action model is updated.
        This update samples random (s, a, r, s') sets from the buffer
        and calculates the discounted reward for each set.
        The value of the actions at states s and s' are predicted from
        the value model. The action model is updated using these value
        predictions as the targets. The value of performed action is
        updated with the discounted reward (using its value prediction
        at s'). ie. x=s, y=[action value 1, action value 2].
        """
        # If buffer isn't full, don't train
        if not self.replay_buffer.full:
            return
        # Else sample batch from buffer
        ss, aa, rr, dd, ss_ = self.replay_buffer.sample_batch(
            self.replay_buffer_samples
        )
        # Calculate estimated S,A values for current states and next states.
        # These are stacked together first to avoid making two separate
        # predict calls (which is slow on GPU).
        ss = np.array(ss)
        ss_ = np.array(ss_)
        y_now_and_future = self._target_model.predict_on_batch(np.vstack((ss, ss_)))
        # Separate again
        y_now = y_now_and_future[0 : self.replay_buffer_samples]
        y_future = y_now_and_future[self.replay_buffer_samples : :]
        # Update rewards where not done with y_future predictions
        dd_mask = np.array(dd, dtype=bool).squeeze()
        rr = np.array(rr, dtype=float).squeeze()
        # Gather max action indexes and update relevant actions in y
        if self.double:
            # If using double dqn select best actions using the action model,
            # but the value of those action using the
            # target model (already have in y_future).
            y_future_action_model = self._action_model.predict_on_batch(ss_)
            selected_actions = np.argmax(y_future_action_model[~dd_mask, :], axis=1)
        else:
            # If normal dqn select targets using target model,
            # and value of those from target model too
            selected_actions = np.argmax(y_future[~dd_mask, :], axis=1)
        # Update reward values with estimated values (where not done)
        # and final rewards (where done)
        rr[~dd_mask] += y_future[~dd_mask, selected_actions]
        if self.final_reward is not None:
            # If self.final_reward is set, set done cases to this value.
            # Else leave as observed reward.
            rr[dd_mask] = self.final_reward
        aa = np.array(aa, dtype=int)
        np.put_along_axis(y_now, aa.reshape(-1, 1), rr.reshape(-1, 1), axis=1)
        # Fit model with updated y_now values
        self._action_model.train_on_batch(ss, y_now)

    def get_best_action(self, s: np.ndarray) -> np.ndarray:
        """
        Get best action(s) from model - the one with the highest predicted value.
        :param s: A single or multiple rows of state observations.
        :return: The selected action.
        """
        preds = self._action_model.predict(self.transform(s))
        return np.argmax(preds)

    def get_action(self, s: np.ndarray, training: bool = False) -> int:
        """
        Get an action using get_best_action or epsilon greedy.
        Epsilon decays every time a random action is chosen.
        :param s: The raw state observation.
        :param training: Bool to indicate whether or not to use this
                         experience to update the model. If False, just
                         returns best action.
        :return: The selected action.
        """
        action = self.eps.select(
            greedy_option=lambda: self.get_best_action(s),
            random_option=lambda: self.env.action_space.sample(),
            training=training,
        )
        return action

    def update_target_model(self) -> None:
        """
        Update the value model with the weights of the action model
        (which is updated each step).
        The value model is updated less often to aid stability.
        """
        self._target_model.set_weights(self._action_model.get_weights())

    def after_episode_update(self) -> None:
        """Value model synced with action model at the end of each episode."""
        self.update_target_model()

    def _discounted_reward(
        self, reward: float, estimated_future_action_rewards: np.ndarray
    ) -> float:
        """
        Use this to define the discounted reward for unfinished episodes,
        default is 1 step TD.
        """
        return reward + self.gamma * np.max(estimated_future_action_rewards)

    def _get_reward(
        self, reward: float, estimated_future_action_rewards: np.ndarray, done: bool
    ) -> float:
        """
        Calculate discounted reward for a single step.
        :param reward: Last real reward.
        :param estimated_future_action_rewards: Estimated future values
                                                of actions taken on next step.
        :param done: Flag indicating if this is the last step on an episode.
        :return: Reward.
        """
        if done:
            # If done, reward is just this step. Can finish because agent has won or lost.
            return self._final_reward(reward)
        else:
            # Otherwise, it's the reward plus the predicted max value of next action
            return self._discounted_reward(reward, estimated_future_action_rewards)


# # Play single episode
# An episode refers to one play through an enviroment. It will take many episodes to train an agent
def play_episode(
    env: gym.Env, agent, n_steps: int = 10, pr: bool = False, training: bool = False
) -> Tuple[List[float], List[int]]:
    episode_rewards = []
    episode_actions = []
    obs = env.reset()
    done = False
    for s in range(n_steps):
        if done:
            break
        # Select action
        action = agent.get_action(obs)
        episode_actions.append(action)
        # Take action
        prev_obs = obs
        obs, reward, done, info = env.step(action)
        episode_rewards.append(reward)
        # Update model
        if training:
            agent.update_experience(s=prev_obs, a=action, r=reward, d=done)
            agent.update_model()
        if pr:
            print(f"Step {s}: Action taken {action}, " f"reward recieved {reward}")
        last_step = s
    if training:
        agent.after_episode_update()
    return episode_rewards, episode_actions


env = gym.make("GFootball-11_vs_11_kaggle-simple115v2-v0")
agent = DeepQAgent(
    # Realistically this should be more like 100,000!
    replay_buffer=ContinuousBuffer(buffer_size=2000),
    eps=EpsilonGreedy(),
    model_architecture=DenseNN(
        observation_shape=env.observation_space.shape,
        # This is also set too low to make running in a notebook easier
        unit_scale=0.3,
        n_actions=env.action_space.n,
    ),
)
_ = play_episode(env, agent, n_steps=100, pr=True, training=True)


#
# # Play multiple episodes
def run_multiple_episodes(
    env: gym.Env,
    agent,
    n_episodes: int = 10,
    n_steps: int = 10,
    pr: bool = False,
    training: bool = False,
) -> List[float]:
    total_episode_rewards = []
    for ep in range(n_episodes):
        episode_rewards, _ = play_episode(env, agent, n_steps, pr=False, training=True)
        total_episode_rewards.append(sum(episode_rewards))
        if pr:
            print(
                f"Episode {ep} finished after {len(episode_rewards)} "
                f"steps, total reward: {sum(episode_rewards)}"
            )
    return total_episode_rewards


env = gym.make("GFootball-11_vs_11_kaggle-simple115v2-v0")
agent = DeepQAgent(
    replay_buffer=ContinuousBuffer(buffer_size=1000),
    eps=EpsilonGreedy(eps_initial=0.7),
    model_architecture=DenseNN(
        observation_shape=env.observation_space.shape,
        unit_scale=0.25,
        n_actions=env.action_space.n,
    ),
)
# Both n_episodes and n_steps need to
reward_history = run_multiple_episodes(
    env,
    agent,
    n_episodes=10,
    # Use 3000 for 90 minute games
    n_steps=3000,
    training=True,
    pr=True,
)
plt.plot(reward_history)
plt.xlabel("Episode")
plt.ylabel("Reward")
plt.show()


# # Compare to a random agent
class RandomModel:
    def __init__(self, n_actions: int):
        self.n_actions = n_actions

    def predict(self) -> int:
        return int(np.random.randint(0, self.n_actions, 1))


@dataclass
class RandomAgent:
    """A random agent that acts randomly and learns nothing."""

    n_actions: int
    name: str = "RandomAgent"

    def __post_init__(self) -> None:
        self._build_model()

    def _build_model(self) -> None:
        """Set model function."""
        self.model = RandomModel(self.n_actions)

    def update_model(self, *args, **kwargs) -> None:
        """No model to update."""
        pass

    def get_action(self, s: Any, **kwargs) -> int:
        return self.model.predict()


env = gym.make("GFootball-11_vs_11_kaggle-simple115v2-v0")
obs = env.reset()
random_agent = RandomAgent(n_actions=env.action_space.n)
ra_reward_history = run_multiple_episodes(
    env, agent, n_episodes=25, n_steps=3000, training=False
)
dqn_reward_history = run_multiple_episodes(
    env, agent, n_episodes=25, n_steps=3000, training=False
)
plt.plot(ra_reward_history, label="Random")
plt.plot(dqn_reward_history, label="DQN")
plt.xlabel("Episode")
plt.ylabel("Reward")
plt.legend(title="Agent")
plt.show()


# # Prepare submission
# This is covered in a bit more detail in the next notebook: https://www.kaggle.com/garethjns/convolutional-deep-q-learner . In short, we need to create a single file containing the agent components we need for inference, including the serialsed model weights. We also need to add to agent function that acts given an environment observation.
# ## Get model weights
import zlib
import pickle

weights_bytes = zlib.compress(pickle.dumps(agent._action_model.get_weights()))

# ## Create submission.py
# Copy and paste the contents of weight_bytes into WEIGHTS_BYTES
import zlib
import pickle
import tensorflow as tf
import numpy as np
from tensorflow import keras
from typing import List, Tuple, Union, Callable
from gfootball.env.wrappers import Simple115StateWrapper


class DenseNN:
    def __init__(
        self,
        observation_shape: List[int],
        n_actions: int,
        output_activation: Union[None, str] = None,
        unit_scale: int = 1,
        learning_rate: float = 0.0001,
        opt: str = "Adam",
    ) -> None:
        """
        :param observation_shape: Tuple specifying input shape.
        :param n_actions: Int specifying number of outputs
        :param output_activation: Activation function for output. Eg.
                                  None for value estimation (off-policy methods).
        :param unit_scale: Multiplier for all units in FC layers in network
                           (not used here at the moment).
        :param opt: Keras optimiser to use. Should be string.
                    This is to avoid storing TF/Keras objects here.
        :param learning_rate: Learning rate for optimiser.
        """
        self.observation_shape = observation_shape
        self.n_actions = n_actions
        self.unit_scale = unit_scale
        self.output_activation = output_activation
        self.learning_rate = learning_rate
        self.opt = opt

    def _model_architecture(self) -> Tuple[keras.layers.Layer, keras.layers.Layer]:
        frame_input = keras.layers.Input(name="input", shape=self.observation_shape)
        fc1 = keras.layers.Dense(
            int(self.observation_shape[0] / 1.5 * self.unit_scale),
            name="fc1",
            activation="relu",
        )(frame_input)
        fc2 = keras.layers.Dense(
            int(self.observation_shape[0] / 3 * self.unit_scale),
            name="fc2",
            activation="relu",
        )(fc1)
        fc3 = keras.layers.Dense(self.n_actions * 2, name="fc3", activation="relu")(fc2)
        action_output = keras.layers.Dense(
            units=self.n_actions, name="output", activation=self.output_activation
        )(fc3)
        return frame_input, action_output

    def compile(
        self, model_name: str = "model", loss: Union[str, Callable] = "mse"
    ) -> keras.Model:
        """
        Compile a copy of the model using the provided loss.
        :param model_name: Name of model
        :param loss: Model loss. Default 'mse'. Can be custom callable.
        """
        # Get optimiser
        if self.opt.lower() == "adam":
            opt = keras.optimizers.Adam
        elif self.opt.lower() == "rmsprop":
            opt = keras.optimizers.RMSprop
        else:
            raise ValueError(f"Invalid optimiser {self.opt}")
        state_input, action_output = self._model_architecture()
        model = keras.Model(
            inputs=[state_input], outputs=[action_output], name=model_name
        )
        model.compile(optimizer=opt(learning_rate=self.learning_rate), loss=loss)
        return model

    def plot(self, model_name: str = "model") -> None:
        keras.utils.plot_model(
            self.compile(model_name), to_file=f"{model_name}.png", show_shapes=True
        )
        plt.show()


tf_mod = DenseNN(
    observation_shape=(115,),
    # This needs to match the number from training
    unit_scale=0.25,
    n_actions=19,
).compile()
tf_mod.set_weights(pickle.loads(zlib.decompress(WEIGHTS_BYTES)))


def agent(obs):
    # Use the existing model and obs buffer on each call to agent
    global tf_mod
    # Convert these to the same output as the Simple115StateWrapper we used in training
    obs = Simple115StateWrapper.convert_observation(
        obs["players_raw"], fixed_positions=False
    )
    # Predict actions from keras model
    actions = tf_mod.predict(obs)
    action = np.argmax(actions)
    return [action]


# # Test submission.py
# The written submission file can be tested with the following block.
# See here: https://github.com/garethjns/kaggle-football/blob/main/scripts/debug_agent.py for a version that can be used to run locally, but also maintain debugability.
from kaggle_environments import make

env = make(
    "football", configuration={"save_video": True, "scenario_name": "11_vs_11_kaggle"}
)
# Define players
left_player = "submission.py"  # A custom agent, eg. random_agent.py or example_agent.py
right_player = "run_right"  # eg. A built in 'AI' agent or the agent again
output = env.run([left_player, right_player])
print(
    f"Final score: {sum([r['reward'] for r in output[0]])} : {sum([r['reward'] for r in output[1]])}"
)
env.render(mode="human", width=800, height=600)

# # Conclusions and next steps
# The DQN performs currently extremely poorly, which is expected given the totally inadequate training done so far here. However, it seems to be running, so that's a start.
# A few things to try:
# - Increase number of training iterations
# - Increase size of replay buffer
# - Increase number of units in network (set unit_scale higher)
# - Increase depth of NN (modify architecture)
# # Appendix
# ## Training with reinforcement-learning-keras
# Follows the same logic as above, but is a bit more fancy. https://github.com/garethjns/reinforcement-learning-keras
import gfootball
from reinforcement_learning_keras.agents.components.history.training_history import (
    TrainingHistory,
)
from reinforcement_learning_keras.agents.components.replay_buffers.continuous_buffer import (
    ContinuousBuffer,
)
from reinforcement_learning_keras.agents.models.dense_nn import DenseNN
from reinforcement_learning_keras.agents.q_learning.deep_q_agent import DeepQAgent
from reinforcement_learning_keras.agents.q_learning.exploration.epsilon_greedy import (
    EpsilonGreedy,
)

agent = DeepQAgent(
    name="deep_q",
    model_architecture=DenseNN(observation_shape=(115,), n_actions=19),
    replay_buffer=ContinuousBuffer(buffer_size=30000),
    env_spec="GFootball-11_vs_11_kaggle-simple115v2-v0",
    eps=EpsilonGreedy(
        eps_initial=0.9, decay=0.0001, eps_min=0.01, decay_schedule="linear"
    ),
    training_history=TrainingHistory(
        agent_name="deep_q", rolling_average=1, plotting_on=True, plot_every=5
    ),
)
agent.train(
    verbose=True,
    render=False,
    n_episodes=10,
    max_episode_steps=3000,
    update_every=2,
    checkpoint_every=1,
)
