# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/05_data.block.ipynb (unless otherwise specified).

__all__ = ['DQN', 'TestDataset']

# Cell
# Python native modules
import os
# Third party libs
from fastcore.all import *
from fastai.torch_basics import *
from fastai.data.all import *
from fastai.basics import *
from torch.utils.data import Dataset
from torch import nn
# Local modules

# Cell
class DQN(Module):
    def __init__(self):
        self.policy=nn.Sequential(
            nn.Linear(4,50),
            nn.ReLU(),
            nn.Linear(50,2)
        )

    def forward(self,x):
        return torch.argmax(self.policy(x),dim=0)

# Cell
class TestDataset(Dataset):
    def __init__(self,policy,device='cpu'):
        self.policy=policy
        self.env=gym.make('CartPole-v1')
        self.next_state=self.env.reset()
        self.device='cpu'

    def __len__(self): return 100
    def __getitem__(self,idx):
        print(id(self.env),' ')
        self.next_state, r, is_done, _=self.env.step(self.policy(Tensor(self.next_state).to(self.device)).numpy())
        if is_done:self.next_state=self.env.reset()
        return self.next_state

# Cell
# Python native modules
import os
# Third party libs
from fastcore.all import *
from fastai.torch_basics import *
from fastai.data.all import *
from fastai.basics import *
from torch.utils.data import Dataset
from torch import nn
# Local modules
from .block import *