# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/10c_agents.dqn.double.ipynb (unless otherwise specified).

__all__ = ['DoubleDQNTrainer']

# Cell
# Python native modules
import os
from collections import deque
# Third party libs
import torch
from torch.nn import *
from fastcore.all import *
from fastai.learner import *
from fastai.torch_basics import *
from fastai.torch_core import *
from fastai.callback.all import *
# Local modules
from ...data.block_simple import *
from ...data.gym import *
from ...agent import *
from ...core import *
from .core import *
from .targets import *

# Cell
class DoubleDQNTrainer(DQNTargetTrainer):
    def after_pred(self):
        self.learn.yb=self.xb
        self.learn.xb=self.xb[0]

        self._xb=({k:v.clone() for k,v in self.xb.items()},)
        self.learn.done_mask=self.xb['done'].reshape(-1,)
        chosen_actions=self.learn.next_q=self.model.model(self.xb['next_state']).argmax(dim=1).reshape(-1,1)
        self.learn.next_q=self.target_model(self.xb['next_state']).gather(1,chosen_actions)
        self.learn.next_q[self.done_mask]=0
        self.learn.targets=self.xb['reward']+self.learn.next_q*(self.discount**self.n_steps)
        self.learn.pred=self.learn.model.model(self.xb['state'])
        t_q=self.pred.clone()
        t_q.scatter_(1,self.xb['action'],self.targets)
        self.learn.yb=(t_q,)