<commit_before>from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import numpy as np
import random
import time
import torch
from torch.autograd import Variable

from optims.helpers import adjust_learning_rate
from core.agent import Agent

class DQNAgent(Agent):
    def __init__(self, args, env_prototype, model_prototype, memory_prototype):
        super(DQNAgent, self).__init__(args, env_prototype, model_prototype, memory_prototype)
        self.logger.warning("<===================================> DQN")

        # env
        self.env = self.env_prototype(self.env_params)
        self.state_shape = self.env.state_shape
        self.action_dim  = self.env.action_dim

        # model
        self.model_params.state_shape = self.state_shape
        self.model_params.action_dim = self.action_dim
        self.model = self.model_prototype(self.model_params)
        self._load_model(self.model_file)   # load pretrained model if provided
        # target_model
        self.target_model = self.model_prototype(self.model_params)
        self._update_target_model_hard()

        # memory
        # NOTE: we instantiate memory objects only inside fit_model/test_model
        # NOTE: since in fit_model we need both replay memory and recent memory
        # NOTE: while in test_model we only need recent memory, in which case memory_size=0
        self.memory_params = args.memory_params

        # experience & states
        self._reset_states()

    def _reset_training_loggings(self):
        self._reset_testing_loggings()
        # during the evaluation in training, we additionally log for
        # the predicted Q-values and TD-errors on validation data
        self.v_avg_log = []
        self.tderr_avg_log = []
        # placeholders for windows for online curve plotting
        if self.visualize:
            self.win_v_avg = "win_v_avg"
            self.win_tderr_avg = "win_tderr_avg"

    def _reset_testing_loggings(self):
        # setup logging for testing/evaluation stats
        self.steps_avg_log = []
        self.steps_std_log = []
        self.reward_avg_log = []
        self.reward_std_log = []
        self.nepisodes_log = []
        self.nepisodes_solved_log = []
        self.repisodes_solved_log = []
        # placeholders for windows for online curve plotting
        if self.visualize:
            self.win_steps_avg = "win_steps_avg"
            self.win_steps_std = "win_steps_std"
            self.win_reward_avg = "win_reward_avg"
            self.win_reward_std = "win_reward_std"
            self.win_nepisodes = "win_nepisodes"
            self.win_nepisodes_solved = "win_nepisodes_solved"
            self.win_repisodes_solved = "win_repisodes_solved"

    def _reset_states(self):
        self._reset_experience()
        self.recent_action = None
        self.recent_observation = None

    # Hard update every `target_model_update` steps.
    def _update_target_model_hard(self):
        self.target_model.load_state_dict(self.model.state_dict())

    # Soft update with `(1 - target_model_update) * old + target_model_update * new`.
    def _update_target_model_soft(self):
        for i, (key, target_weights) in enumerate(self.target_model.state_dict().iteritems()):
            target_weights += self.target_model_update * self.model.state_dict()[key]

    def _sample_validation_data(self):
        self.logger.warning("Validation Data @ Step: " + str(self.step))
        self.valid_data = self.memory.sample(self.valid_size)

    def _compute_validation_stats(self):
        return self._get_q_update(self.valid_data)

    def _get_q_update(self, experiences): # compute temporal difference error for a batch
        # Start by extracting the necessary parameters (we use a vectorized implementation).
        state0_batch_vb    = Variable(torch.from_numpy(np.array(tuple(experiences[i].state0 for i in range(len(experiences))))).type(self.dtype))
        action_batch_vb    = Variable(torch.from_numpy(np.array(tuple(experiences[i].action for i in range(len(experiences))))).long())
        reward_batch_vb    = Variable(torch.from_numpy(np.array(tuple(experiences[i].reward for i in range(len(experiences)))))).type(self.dtype)
        state1_batch_vb    = Variable(torch.from_numpy(np.array(tuple(experiences[i].state1 for i in range(len(experiences))))).type(self.dtype))
        terminal1_batch_vb = Variable(torch.from_numpy(np.array(tuple(0. if experiences[i].terminal1 else 1. for i in range(len(experiences)))))).type(self.dtype)

        if self.use_cuda:
            action_batch_vb = action_batch_vb.cuda()

        # Compute target Q values for mini-batch update.
        if self.enable_double_dqn:
            # According to the paper "Deep Reinforcement Learning with Double Q-learning"
            # (<NAME> al., 2015), in Double DQN, the online network predicts the actions
            # while the target network is used to estimate the Q value.
            q_values_vb = self.model(state1_batch_vb)
            # Detach this variable from the current graph since we don't want gradients to propagate
            q_values_vb = Variable(q_values_vb.data)
            # _, q_max_actions_vb = q_values_vb.max(dim=1)              # 0.1.12
            _, q_max_actions_vb = q_values_vb.max(dim=1, keepdim=True)  # 0.2.0
            # Now, estimate Q values using the target network but select the values with the
            # highest Q value wrt to the online model (as computed above).
            next_max_q_values_vb = self.target_model(state1_batch_vb)
            # Detach this variable from the current graph since we don't want gradients to propagate
            next_max_q_values_vb = Variable(next_max_q_values_vb.data)
            next_max_q_values_vb = next_max_q_values_vb.gather(1, q_max_actions_vb)
        else:
            # Compute the q_values given state1, and extract the maximum for each sample in the batch.
            # We perform this prediction on the target_model instead of the model for reasons
            # outlined in Mnih (2015). In short: it makes the algorithm more stable.
            next_max_q_values_vb = self.target_model(state1_batch_vb)
            # Detach this variable from the current graph since we don't want gradients to propagate
            next_max_q_values_vb = Variable(next_max_q_values_vb.data)
            # next_max_q_values_vb, _ = next_max_q_values_vb.max(dim = 1)               # 0.1.12
            next_max_q_values_vb, _ = next_max_q_values_vb.max(dim = 1, keepdim=True)   # 0.2.0

        # Compute r_t + gamma * max_a Q(s_t+1, a) and update the targets accordingly
        # but only for the affected output units (as given by action_batch).
        current_q_values_vb = self.model(state0_batch_vb).gather(1, action_batch_vb.unsqueeze(1)).squeeze()
        # Set discounted reward to zero for all states that were terminal.
        next_max_q_values_vb = next_max_q_values_vb * terminal1_batch_vb.unsqueeze(1)
        # expected_q_values_vb = reward_batch_vb + self.gamma * next_max_q_values_vb            # 0.1.12
        expected_q_values_vb = reward_batch_vb + self.gamma * next_max_q_values_vb.squeeze()    # 0.2.0
        # Compute temporal difference error, use huber loss to mitigate outlier impact
        # TODO: can optionally use huber loss from here: https://medium.com/@karpathy/yes-you-should-understand-backprop-e2f06eab496b
        td_error_vb = self.value_criteria(current_q_values_vb, expected_q_values_vb)

        # return v_avg, tderr_avg_vb
        if not self.training:   # then is being called from _compute_validation_stats, which is just doing inference
            td_error_vb = Variable(td_error_vb.data) # detach it from the graph
        return next_max_q_values_vb.data.clone().mean(), td_error_vb

    def _epsilon_greedy(self, q_values_ts):
        # calculate epsilon
        if self.training:   # linearly anneal epsilon
            self.eps = self.eps_end + max(0, (self.eps_start - self.eps_end) * (self.eps_decay - max(0, self.step - self.learn_start)) / self.eps_decay)
        else:
            self.eps = self.eps_eval
        # choose action
        if np.random.uniform() < self.eps:  # then we choose a random action
            action = random.randrange(self.action_dim)
        else:                               # then we choose the greedy action
            if self.use_cuda:
                action = np.argmax(q_values_ts.cpu().numpy())
            else:
                action = np.argmax(q_values_ts.numpy())
        return action

    def _forward(self, observation):
        # Select an action.
        state = self.memory.get_recent_state(observation)
        state_ts = torch.from_numpy(np.array(state)).unsqueeze(0).type(self.dtype)
        q_values_ts = self.model(Variable(state_ts, volatile=True)).data # NOTE: only doing inference here, so volatile=True
        if self.training and self.step < self.learn_start:  # then we don't do any learning, just accumulate experiences into replay memory
            action = random.randrange(self.action_dim)      # thus we only randomly sample actions here, since the model hasn't been updated at all till now
        else:
            action = self._epsilon_greedy(q_values_ts)

        # Book keeping
        self.recent_observation = observation
        self.recent_action = action

        return action

    def _backward(self, reward, terminal):
        # Store most recent experience in memory.
        if self.step % self.memory_interval == 0:
            # NOTE: so the tuples stored in memory corresponds to:
            # NOTE: in recent_observation(state0), take recent_action(action), get reward(reward), ends up in terminal(terminal1)
            self.memory.append(self.recent_observation, self.recent_action, reward, terminal,
                               training = self.training)

        if not self.training:
            # We're done here. No need to update the replay memory since we only use the
            # recent memory to obtain the state over the most recent observations.
            return

        # sample validation data right before training started
        # NOTE: here validation data is not entirely clean since the agent might see those data during training
        # NOTE: but it's ok as is also the case in the original dqn code, cos those data are not used to judge performance like in supervised learning
        # NOTE: but rather to inspect the whole learning procedure; of course we can separate those entirely from the training data but it's not worth the effort
        if self.step == self.learn_start + 1:
            self._sample_validation_data()
            self.logger.warning("Start  Training @ Step: " + str(self.step))

        # Train the network on a single stochastic batch.<commit_msg>Modified the DQN agent so that it can produce a list of actions (useful for continuous action spaces). Tested with the opensim-rl environment which has a continous action space<commit_after>from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import numpy as np
import random
import time
import torch
from torch.autograd import Variable

from optims.helpers import adjust_learning_rate
from core.agent import Agent

class DQNAgent(Agent):
    def __init__(self, args, env_prototype, model_prototype, memory_prototype):
        super(DQNAgent, self).__init__(args, env_prototype, model_prototype, memory_prototype)
        self.logger.warning("<===================================> DQN")

        # env
        self.env = self.env_prototype(self.env_params)
        self.state_shape = self.env.state_shape
        self.action_dim  = self.env.action_dim

        # model
        self.model_params.state_shape = self.state_shape
        self.model_params.action_dim = self.action_dim
        self.model = self.model_prototype(self.model_params)
        self._load_model(self.model_file)   # load pretrained model if provided
        # target_model
        self.target_model = self.model_prototype(self.model_params)
        self._update_target_model_hard()

        # memory
        # NOTE: we instantiate memory objects only inside fit_model/test_model
        # NOTE: since in fit_model we need both replay memory and recent memory
        # NOTE: while in test_model we only need recent memory, in which case memory_size=0
        self.memory_params = args.memory_params

        # experience & states
        self._reset_states()

    def _reset_training_loggings(self):
        self._reset_testing_loggings()
        # during the evaluation in training, we additionally log for
        # the predicted Q-values and TD-errors on validation data
        self.v_avg_log = []
        self.tderr_avg_log = []
        # placeholders for windows for online curve plotting
        if self.visualize:
            self.win_v_avg = "win_v_avg"
            self.win_tderr_avg = "win_tderr_avg"

    def _reset_testing_loggings(self):
        # setup logging for testing/evaluation stats
        self.steps_avg_log = []
        self.steps_std_log = []
        self.reward_avg_log = []
        self.reward_std_log = []
        self.nepisodes_log = []
        self.nepisodes_solved_log = []
        self.repisodes_solved_log = []
        # placeholders for windows for online curve plotting
        if self.visualize:
            self.win_steps_avg = "win_steps_avg"
            self.win_steps_std = "win_steps_std"
            self.win_reward_avg = "win_reward_avg"
            self.win_reward_std = "win_reward_std"
            self.win_nepisodes = "win_nepisodes"
            self.win_nepisodes_solved = "win_nepisodes_solved"
            self.win_repisodes_solved = "win_repisodes_solved"

    def _reset_states(self):
        self._reset_experience()
        self.recent_action = None
        self.recent_observation = None

    # Hard update every `target_model_update` steps.
    def _update_target_model_hard(self):
        self.target_model.load_state_dict(self.model.state_dict())

    # Soft update with `(1 - target_model_update) * old + target_model_update * new`.
    def _update_target_model_soft(self):
        for i, (key, target_weights) in enumerate(self.target_model.state_dict().iteritems()):
            target_weights += self.target_model_update * self.model.state_dict()[key]

    def _sample_validation_data(self):
        self.logger.warning("Validation Data @ Step: " + str(self.step))
        self.valid_data = self.memory.sample(self.valid_size)

    def _compute_validation_stats(self):
        return self._get_q_update(self.valid_data)

    def _get_q_update(self, experiences): # compute temporal difference error for a batch
        # Start by extracting the necessary parameters (we use a vectorized implementation).
        state0_batch_vb    = Variable(torch.from_numpy(np.array(tuple(experiences[i].state0 for i in range(len(experiences))))).type(self.dtype))
        action_batch_vb    = Variable(torch.from_numpy(np.array(tuple(experiences[i].action for i in range(len(experiences))))).long())
        reward_batch_vb    = Variable(torch.from_numpy(np.array(tuple(experiences[i].reward for i in range(len(experiences)))))).type(self.dtype)
        state1_batch_vb    = Variable(torch.from_numpy(np.array(tuple(experiences[i].state1 for i in range(len(experiences))))).type(self.dtype))
        terminal1_batch_vb = Variable(torch.from_numpy(np.array(tuple(0. if experiences[i].terminal1 else 1. for i in range(len(experiences)))))).type(self.dtype)

        if self.use_cuda:
            action_batch_vb = action_batch_vb.cuda()

        # Compute target Q values for mini-batch update.
        if self.enable_double_dqn:
            # According to the paper "Deep Reinforcement Learning with Double Q-learning"
            # (<NAME> al., 2015), in Double DQN, the online network predicts the actions
            # while the target network is used to estimate the Q value.
            q_values_vb = self.model(state1_batch_vb)
            # Detach this variable from the current graph since we don't want gradients to propagate
            q_values_vb = Variable(q_values_vb.data)
            # _, q_max_actions_vb = q_values_vb.max(dim=1)              # 0.1.12
            _, q_max_actions_vb = q_values_vb.max(dim=1, keepdim=True)  # 0.2.0
            # Now, estimate Q values using the target network but select the values with the
            # highest Q value wrt to the online model (as computed above).
            next_max_q_values_vb = self.target_model(state1_batch_vb)
            # Detach this variable from the current graph since we don't want gradients to propagate
            next_max_q_values_vb = Variable(next_max_q_values_vb.data)
            next_max_q_values_vb = next_max_q_values_vb.gather(1, q_max_actions_vb)
        else:
            # Compute the q_values given state1, and extract the maximum for each sample in the batch.
            # We perform this prediction on the target_model instead of the model for reasons
            # outlined in Mnih (2015). In short: it makes the algorithm more stable.
            next_max_q_values_vb = self.target_model(state1_batch_vb)
            # Detach this variable from the current graph since we don't want gradients to propagate
            next_max_q_values_vb = Variable(next_max_q_values_vb.data)
            # next_max_q_values_vb, _ = next_max_q_values_vb.max(dim = 1)               # 0.1.12
            next_max_q_values_vb, _ = next_max_q_values_vb.max(dim = 1, keepdim=True)   # 0.2.0

        # Compute r_t + gamma * max_a Q(s_t+1, a) and update the targets accordingly
        # but only for the affected output units (as given by action_batch).
        current_q_values_vb = self.model(state0_batch_vb).gather(1, action_batch_vb.unsqueeze(1)).squeeze()
        # Set discounted reward to zero for all states that were terminal.
        next_max_q_values_vb = next_max_q_values_vb * terminal1_batch_vb.unsqueeze(1)
        # expected_q_values_vb = reward_batch_vb + self.gamma * next_max_q_values_vb            # 0.1.12
        expected_q_values_vb = reward_batch_vb + self.gamma * next_max_q_values_vb.squeeze()    # 0.2.0
        # Compute temporal difference error, use huber loss to mitigate outlier impact
        # TODO: can optionally use huber loss from here: https://medium.com/@karpathy/yes-you-should-understand-backprop-e2f06eab496b
        td_error_vb = self.value_criteria(current_q_values_vb, expected_q_values_vb)

        # return v_avg, tderr_avg_vb
        if not self.training:   # then is being called from _compute_validation_stats, which is just doing inference
            td_error_vb = Variable(td_error_vb.data) # detach it from the graph
        return next_max_q_values_vb.data.clone().mean(), td_error_vb

    def _epsilon_greedy(self, q_values_ts):
        # calculate epsilon
        if self.training:   # linearly anneal epsilon
            self.eps = self.eps_end + max(0, (self.eps_start - self.eps_end) * (self.eps_decay - max(0, self.step - self.learn_start)) / self.eps_decay)
        else:
            self.eps = self.eps_eval
        # choose action
        if np.random.uniform() < self.eps:  # then we choose a random action
            action = np.random.rand(self.action_dim).tolist()
        else:                               # then we choose the greedy action
            if self.use_cuda:
                action = np.argmax(q_values_ts.cpu().numpy())
            else:
                action = np.argmax(q_values_ts.numpy())
        return action

    def _forward(self, observation):
        # Select an action.
        state = self.memory.get_recent_state(observation)
        state_ts = torch.from_numpy(np.array(state)).unsqueeze(0).type(self.dtype)
        q_values_ts = self.model(Variable(state_ts, volatile=True)).data # NOTE: only doing inference here, so volatile=True
        if self.training and self.step < self.learn_start:  # then we don't do any learning, just accumulate experiences into replay memory
            action = np.random.rand(self.action_dim).tolist()      # thus we only randomly sample actions here, since the model hasn't been updated at all till now
        else:
            action = self._epsilon_greedy(q_values_ts)

        # Book keeping
        self.recent_observation = observation
        self.recent_action = action

        return action

    def _backward(self, reward, terminal):
        # Store most recent experience in memory.
        if self.step % self.memory_interval == 0:
            # NOTE: so the tuples stored in memory corresponds to:
            # NOTE: in recent_observation(state0), take recent_action(action), get reward(reward), ends up in terminal(terminal1)
            self.memory.append(self.recent_observation, self.recent_action, reward, terminal,
                               training = self.training)

        if not self.training:
            # We're done here. No need to update the replay memory since we only use the
            # recent memory to obtain the state over the most recent observations.
            return

        # sample validation data right before training started
        # NOTE: here validation data is not entirely clean since the agent might see those data during training
        # NOTE: but it's ok as is also the case in the original dqn code, cos those data are not used to judge performance like in supervised learning
        # NOTE: but rather to inspect the whole learning procedure; of course we can separate those entirely from the training data but it's not worth the effort
        if self.step == self.learn_start + 1:
            self._sample_validation_data()
            self.logger.warning("Start  Training @ Step: " + str(self.step))

        # Train the network on a single stochastic batch.
{'fraction_non_alphanumeric': 0.0797475879098267, 'fraction_numerical': 0.006127394942612831, 'mean_word_length': 2.8041398504087667, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 7, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '305374', 'n_tokens_mistral': 6484, 'n_tokens_neox': 6260, 'n_words': 1731}