text: # SPDX-License-Identifier: Apache-2.0
  
  from collections import deque
  from keras.models import Sequential
  from keras.layers import Dense, Conv2D, MaxPool2D, Flatten
  from copy import deepcopy
  import random
  from keras.optimizers import Adam
  from keras import backend as K
  import tensorflow as tf
  import numpy as np
  
  def reshape_state(state, is_atari_env, state_size):
      reshaped = state
      if not is_atari_env:
          reshaped = np.reshape(state, [1, state_size])
      else:
          if len(state.shape) < 4:
              reshaped = np.expand_dims(state, axis=0)
      return reshaped
  
  def update_loss(loss, sample_loss):
      if loss is not None and sample_loss is not None:
          for key, val in sample_loss.items():
              if key in loss:
                  loss[key] += val
              else:
                  loss[key] = val
  
  def concatenate_state_action(state, action):
      out = np.concatenate((state[0], [action]))
      out = np.reshape(out, [1, len(out)])
      return out
  
  class DQNAgent:
      def __init__(self, state_size, action_size, is_atari_env, is_delayed_agent=False, delay_value=0, epsilon_min=0.001,
                  epsilon_decay=0.999, learning_rate=0.001, epsilon=1.0, use_m_step_reward=False, use_latest_reward=True,
                   loss='mse', **kwargs):
          self.state_size = state_size
          self.action_size = action_size
          self.is_atari_env = is_atari_env
          mem_len = 50000 if self.is_atari_env else 2000
          self.memory = deque(maxlen=mem_len)
          self.gamma = 0.95    # discount rate
          self.epsilon = epsilon  # exploration rate
          self.epsilon_min = epsilon_min
          self.epsilon_decay = epsilon_decay #0.995
          self.learning_rate = learning_rate
          self.sample_buffer = deque()
          self.is_delayed_agent = is_delayed_agent
          self.delay_value = delay_value
          self.model = self._build_model(loss=loss)
          self.use_m_step_reward = use_m_step_reward
          self.use_latest_reward = use_latest_reward
  
  
      def _huber_loss(self, y_true, y_pred, clip_delta=1.0):
          """Huber loss for Q Learning
          References: https://en.wikipedia.org/wiki/Huber_loss
                      https://www.tensorflow.org/api_docs/python/tf/losses/huber_loss
          """
          error = y_true - y_pred
          cond = K.abs(error) <= clip_delta
  
          squared_loss = 0.5 * K.square(error)
          quadratic_loss = 0.5 * K.square(clip_delta) + clip_delta * (K.abs(error) - clip_delta)
  
          return K.mean(tf.where(cond, squared_loss, quadratic_loss))
  
      def _build_model(self, loss=None, input_size=None, output_size=None):
          loss = self._huber_loss if loss is 'huber' else loss
          input_size = self.state_size if input_size is None else input_size
          output_size = self.action_size if output_size is None else output_size
  
          # Neural Net for Deep-Q learning Model
          model = Sequential()
          if self.is_atari_env:
              model.add(Conv2D(32, 8, strides=(4,4), input_shape=input_size, activation='relu'))
              model.add(MaxPool2D())
              model.add(Conv2D(64, 4, strides=(2,2), activation='relu'))
              model.add(MaxPool2D())
              model.add(Conv2D(64, 3, strides=(1,1), activation='relu'))
              model.add(MaxPool2D())
              model.add(Flatten())
              model.add(Dense(64, activation='relu'))
              model.add(Dense(64, activation='relu'))
              model.add(Dense(output_size, activation='linear'))
          else:
              model.add(Dense(24, input_dim=input_size, activation='relu'))
              model.add(Dense(24, activation='relu'))
              model.add(Dense(output_size, activation='linear'))
  
          model.compile(loss=loss,
                        optimizer=Adam(lr=self.learning_rate))
          return model
  
      def memorize(self, state, action, reward, next_state, done):
          if self.is_delayed_agent:
              # for earlier time than delay_value, the data is problematic (non-delayed response)
              # Construct modified tuple by keeping old s_t with new a_{t+m}, r_{t+m} s_{t+m+1}
              new_tuple = (state, action, reward, next_state, done)
              self.sample_buffer.append(new_tuple)
              if len(self.sample_buffer) - 1 >= self.delay_value:
                  old_tuple = self.sample_buffer.popleft()
                  modified_tuple = list(deepcopy(old_tuple))
                  modified_tuple[1] = action
                  modified_tuple[2] = self.m_step_reward(first_reward=old_tuple[2])
                  # trying to use s_{t+1} instead of s_{t+m} as in the original ICML2020 submission
                  # modified_tuple[3] = next_state
                  modified_tuple = tuple(modified_tuple)
                  self.memory.append(modified_tuple)
          else:
              self.memory.append((state, action, reward, next_state, done))
  
      def act(self, state, eval=False):
          if not eval and np.random.rand() <= self.epsilon:
              return random.randrange(self.action_size)
          act_values = self.model.predict(state)
          return np.argmax(act_values[0])  # returns action
  
      def m_step_reward(self, first_reward):
          if not self.use_m_step_reward:
              if self.use_latest_reward:
                  return self.sample_buffer[-1][2]
              else:
                  return first_reward
          else:
              discounted_rew = first_reward
              for i in range(self.delay_value):
                  discounted_rew += self.gamma ** (i + 1) * self.sample_buffer[i][2]
              return discounted_rew
  
      def effective_gamma(self):
          return self.gamma if not self.use_m_step_reward else (self.gamma ** (self.delay_value + 1))
  
      def replay(self, batch_size):
          minibatch = random.sample(self.memory, batch_size)
          for state, action, reward, next_state, done in minibatch:
              target = reward
              if not done:
                  target = (reward + self.effective_gamma() *
                            np.amax(self.model.predict(next_state)[0]))
              target_f = self.model.predict(state)
              target_f[0][action] = target
              # self.model.fit(state, target_f, epochs=1, verbose=0,
              #                callbacks=[WandbCallback()])
              self.model.fit(state, target_f, epochs=1, verbose=0)
          if self.epsilon > self.epsilon_min:
              self.epsilon *= self.epsilon_decay
  
  
      def load(self, name):
          self.model.load_weights(name)
  
      def save(self, name):
          self.model.save_weights(name)
  
      def clear_action_buffer(self):
          self.sample_buffer.clear()
  
  
  class DDQNAgent(DQNAgent):
      def __init__(self, state_size, action_size, is_atari_env, is_delayed_agent=False, delay_value=0, epsilon_min=0.001,
                  epsilon_decay=0.999, learning_rate=0.001, epsilon=1.0, use_m_step_reward=False, use_latest_reward=True):
          super().__init__(state_size, action_size, is_atari_env=is_atari_env, is_delayed_agent=is_delayed_agent, delay_value=delay_value,
                           epsilon_min=epsilon_min, epsilon_decay=epsilon_decay, learning_rate=learning_rate,
                           epsilon=epsilon, use_m_step_reward=use_m_step_reward, use_latest_reward=use_latest_reward,
                           loss='huber')
          # self.model = self._build_model()
          self.target_model = self._build_model(loss='huber')
          self.update_target_model()
  
  
      def update_target_model(self):
          # copy weights from model to target_model
          self.target_model.set_weights(self.model.get_weights())
  
      def train_model(self, batch):
          state_vec, action_vec, reward_vec, next_state_vec, done_vec = batch
          target = self.model.predict(state_vec)
          # a = self.model.predict(next_state)[0]
          t = self.target_model.predict(next_state_vec)#[0]
          not_done_arr = np.invert(np.asarray(done_vec))
          new_targets = reward_vec + not_done_arr * self.effective_gamma() * np.amax(t, axis=1)
          for i in range(len(batch[0])):
              target[i][action_vec[i]] = new_targets[i]
          # target[0][action] = reward + self.gamma * t[np.argmax(a)]
          train_history = self.model.fit(state_vec, target, epochs=1, verbose=0)
          q_loss = train_history.history['loss'][0]
          loss_dict = {'q_loss': q_loss}
          return loss_dict
  
      def _create_batch(self, indices):
          state_vec, action_vec, reward_vec, next_state_vec, done_vec = [], [], [], [], []
          for i in indices:
              data = self.memory[i]
              state, action, reward, next_state, done = data
              state_vec.append(np.array(state, copy=False))
              action_vec.append(action)
              reward_vec.append(reward)
              next_state_vec.append(np.array(next_state, copy=False))
              done_vec.append(done)
          return np.concatenate(state_vec, axis=0), action_vec, reward_vec, np.concatenate(next_state_vec, axis=0), done_vec
  
      def replay(self, batch_size):
          loss = {}
          indices = np.random.choice(len(self.memory), batch_size)
          batch = self._create_batch(indices)
          sample_loss = self.train_model(batch)
          update_loss(loss, sample_loss)
          if self.epsilon > self.epsilon_min:
              self.epsilon *= self.epsilon_decay
          return loss
  
  class DDQNPlanningAgent(DDQNAgent):
      def __init__(self, state_size, action_size, is_atari_env, is_delayed_agent=False, delay_value=0, epsilon_min=0.001,
                   epsilon_decay=0.999, learning_rate=0.001, epsilon=1.0, use_m_step_reward=False,
                   use_latest_reward=True, env=None, use_learned_forward_model=True):
          super().__init__(state_size, action_size, is_atari_env=is_atari_env, is_delayed_agent=is_delayed_agent, delay_value=delay_value,
                           epsilon_min=epsilon_min, epsilon_decay=epsilon_decay, learning_rate=learning_rate,
                           epsilon=epsilon, use_m_step_reward=use_m_step_reward, use_latest_reward=use_latest_reward)
          self.use_learned_forward_model = use_learned_forward_model
          if self.use_learned_forward_model:
              keras_forward_model = self._build_model(loss='mse', input_size=self.state_size + 1, output_size=self.state_size)
              self.forward_model = ForwardModel(keras_forward_model)
          else:
              self.forward_model = env
  
      def train_model(self, batch):
          loss_dict = super().train_model(batch)
          if self.use_learned_forward_model and self.delay_value > 0:
              state_vec, action_vec, _, next_state_vec, _ = batch
              act_t = np.asarray([action_vec]).transpose()
              concat_vec = np.concatenate((state_vec, act_t), axis=1)
              train_history = self.forward_model.keras_model.fit(concat_vec, next_state_vec, epochs=1, verbose=0)
              f_model_loss = train_history.history['loss'][0]
              loss_dict['f_model_loss'] = f_model_loss
          return loss_dict
  
      def act(self, state, pending_actions, eval):
          if not eval and np.random.rand() <= self.epsilon:
              return random.randrange(self.action_size)
          last_state = state
          if self.delay_value > 0:
              if not self.use_learned_forward_model:
                  self.forward_model.store_initial_state()
                  # initial_state = deepcopy(state)
              for curr_action in pending_actions:
                  last_state = self.forward_model.get_next_state(state=last_state, action=curr_action)
              if not self.use_learned_forward_model:
                  self.forward_model.restore_initial_state()
          last_state_r = reshape_state(last_state, self.is_atari_env, self.state_size)
          act_values = self.model.predict(last_state_r)
          return np.argmax(act_values[0])  # returns best action for last state
  
      def memorize(self, state, action, reward, next_state, done):
          # for earlier time than delay_value, the data is problematic (non-delayed response)
          # Construct modified tuple by keeping old s_t with new a_{t+m}, r_{t+m} s_{t+m+1}
          new_tuple = (state, action, reward, next_state, done)
          self.sample_buffer.append(new_tuple)
          if len(self.sample_buffer) - 1 >= self.delay_value:
              old_tuple = self.sample_buffer.popleft()
              modified_tuple = list(deepcopy(old_tuple))
              # build time-coherent tuple from new tuple and old action
              modified_tuple[0] = state
              # modified_tuple[1] = action
              modified_tuple[2] = reward #self.m_step_reward(first_reward=old_tuple[2])
              modified_tuple[3] = next_state
              modified_tuple = tuple(modified_tuple)
              self.memory.append(modified_tuple)
  
  class ForwardModel:
      def __init__(self, keras_model):
          self.keras_model = keras_model
  
      def get_next_state(self, state, action):
          input = concatenate_state_action(state, action)
          return self.keras_model.predict(input)
  
      def reset_to_state(self, state):
          # not necessary here. Only used if the forwrad_model is the actual env instance
          pass