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

__all__ = ['DQNTargetTrainer']

# Cell
# Python native modules
import os
from collections import deque
from typing import *
# 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 import *
from ...data.gym import *
from ...agent import *
from ...core import *
from .core import *
from ...memory.experience_replay import *

# Cell
class DQNTargetTrainer(Callback):

    def __init__(self,n_batch=0,target_sync=300,discount=0.99,n_steps=1):
        store_attr()
        self._xb=None

    def before_fit(self):
        self.learn.target_model=deepcopy(self.learn.model.model)
        self.n_batch=0

    def after_pred(self):
#         print(self.xb)
        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,)
        self.learn.next_q=self.target_model(self.xb['next_state']).max(dim=1).values.reshape(-1,1)
        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,)

    def before_backward(self): self.learn.xb=self._xb

    def after_batch(self):
        if self.n_batch%self.target_sync==0:
            self.target_model.load_state_dict(self.learn.model.state_dict())
        self.n_batch+=1