<jupyter_start><jupyter_text>Env class<jupyter_code>def apply_aug(p0, aug):
    if aug == 0:
        p = p0.copy()
    elif aug == 1:
        p = cv2.rotate(p0, cv2.ROTATE_90_CLOCKWISE)
    elif aug == 2:
        p = cv2.rotate(p0, cv2.ROTATE_180)
    elif aug == 3:
        p = cv2.rotate(p0, cv2.ROTATE_90_COUNTERCLOCKWISE)
    elif aug == 4:
        p = cv2.flip(p0, 1)
    elif aug == 5:
        p = cv2.rotate(p0, cv2.ROTATE_90_CLOCKWISE)
        p = cv2.flip(p, 1)
    elif aug == 6:
        p = cv2.rotate(p0, cv2.ROTATE_180)
        p = cv2.flip(p, 1)
    elif aug == 7:
        p = cv2.rotate(p0, cv2.ROTATE_90_COUNTERCLOCKWISE)
        p = cv2.flip(p, 1)
    return p

class NeuralEnv:
    def __init__(self,
                 env_model_path: str,
                 reward_model_path: str,
                 device: str,
                 batch_size = 16,
                 reward_confidence=0.5,
                 stone_frac=0.0,
                 step_size=4,
                 max_step=14):
        '''
        input params:
            env_model_path    [str] : path to model s_next=model(s_curr,action)
            reward_model_path [str] : path to model reward=model(s_curr)
            device            [str] : one of {'cpu', 'cuda:0', 'cuda:1'}
            batch_size        [int] : len of batch
            reward_confidence [flt] : classificator's confidence
            stone_frac        [flt] : part of the initial states with guaranteed stones
            step_size         [int] : 
            max_step          [int] : 
        output params:
            all output-variables will be torch.tensors in the selected DEVICE
            all input-variables have to be torch.tensors in the selected DEVICE
        '''
        self.device = device
        self.batch_size = batch_size
        self.reward_confidence = reward_confidence
        self.stone_frac = stone_frac
        self.step_size = step_size
        self.max_step = max_step
        self.reward_frame_transform = transforms.Compose([transforms.CenterCrop(24)])
        self.frame_transform = A.Compose([A.Normalize(mean=(0.5,), std=(0.5,)),
                                          ToTensorV2(transpose_mask=False)])
        
        self.model = ResNetUNet_v2(3)
        self.model.load_state_dict(torch.load(env_model_path))
        self.model = self.model.to(self.device)
        self.model.eval()
        
        self.stone_classifier = StoneClassifier()
        self.stone_classifier.load_state_dict(torch.load(reward_model_path))
        self.stone_classifier = self.stone_classifier.to(self.device)
        self.stone_classifier.eval()

        self.df = pd.read_csv('/media/laggg/surviv_rl_data/dataset_inventory_v2.csv')
        self.df = self.df[self.df.zoom == 1].reset_index()
    #----------------------------------------------------------------------------------------------------
        
    def reset(self):
        '''
        output params:
            init_s     [float torch tensor [-1...1]] : batch of initial states (batch,3,96,96)
            init_supp  [float torch tensor]          : batch of initial support vector (batch,2)
        '''
        init_s = torch.zeros(self.batch_size,3,96,96).float()
        init_supp = torch.zeros(self.batch_size,2).float()
        
        for i in range(self.batch_size):
            j = np.random.randint(len(self.df))
            frame = load_image(self.df["video"][j], self.df["frame"][j])
            aug = np.random.choice(np.arange(8), p=np.array([0.125]*8))
            frame = apply_aug(frame, aug)
            frame = self.frame_transform(image=frame)['image']
            supp = torch.tensor([self.df["sp"][j]/100,self.df["zoom"][j]/15]).float()
            reward = self.get_reward(frame.unsqueeze(0).to(self.device))
            while reward>0:
                frame = frame.unsqueeze(0).expand(8,3,96,96).to(self.device)
                supp = supp.unsqueeze(0).expand(8,2).to(self.device)
                act = torch.tensor([x for x in range(8)]).to(self.device)
                frame, supp, reward = self.step(frame,supp,act)
                k = torch.argmin(reward)
                reward = reward[k]
                frame = frame[k].detach().cpu()
                supp = supp[k].detach().cpu()
            # +условие, чтобы агент c вер-тью p ресался на кадре, на котором где-то есть камень
            init_s[i] = frame
            init_supp[i] = supp
        return init_s.to(self.device),init_supp.to(self.device)
    #----------------------------------------------------------------------------------------------------
    
    def get_reward(self, state):
        '''
        input params:
            state [float torch.tensor [-1...1]] : batch of states (batch,3,96,96)
        output params:
            r      [float torch.tensor [0...1]]  : batch of rewards (batch,1)
        '''
        state = self.reward_frame_transform(state)
        with torch.no_grad():
            r = self.stone_classifier(state)[:,1].unsqueeze(1)
        r = (r>self.reward_confidence).float().detach()
        return r
    #----------------------------------------------------------------------------------------------------
    
    def step(self, s_curr, supp_curr, action):
        '''
        input params:
            s_curr    [float torch.tensor [-1...1]] : batch of current states (batch,3,96,96)
            supp_curr [float torch tensor]          : batch of current support vector (batch,2)
            action    [int torch tensor {0,...,7}]  : batch of chosen direction (batch,)  
        output params:
            s_next    [float torch.tensor [-1...1]] : batch of next states (batch,3,96,96)
            supp_next [float torch tensor]          : batch of next support vector =supp_curr (batch,2)
            reward    [float torch.tensor [0...1]]  : batch of rewards (batch,1)
        '''
        action_ohe = F.one_hot(action, num_classes=8).float() # (batch,8)
        if len(action)==1:
            action_ohe.unsqueeze(0)
        n =  torch.tensor([self.step_size/self.max_step]*action.size()[0])
        n = n.unsqueeze(1).float().to(self.device) # (batch,1)
        v = torch.cat([action_ohe,supp_curr,n], dim=1) # (batch,8+2+1)
        with torch.no_grad():
            s_next = self.model((s_curr,v)).detach()
        reward = self.get_reward(s_next)
        return s_next, supp_curr, reward
def inference_agent(model, s_init, supp_init):
    '''
    s_init (3,96,96) cuda
    supp_init (2,) cuda
    '''
    model.eval()
    s_curr = s_init.unsqueeze(0).clone()
    supp_curr = supp_init.unsqueeze(0).clone()
    
    for i in range(20):
        s_vis = (s_curr[0].permute(1,2,0).detach().cpu()+1)/2
        display.clear_output(wait=True)
        see_plot(s_vis)
        with torch.no_grad():
            act = torch.argmax(model(s_curr), dim=1)
        s_curr, supp_curr, reward = env.step(s_curr, supp_curr, act)
        time.sleep(0.1)<jupyter_output><empty_output><jupyter_text>Example<jupyter_code>params = {'DEVICE': 'cuda:0',
          'BATCH': 128,
         'reward_confidence': 0.95}

# 0. build environment
env = NeuralEnv("./model_weights/resunet_v5.pth",
                "./model_weights/nostone_stone_classifier.pth",
                params['DEVICE'],
                params['BATCH'],
                params['reward_confidence'])

# 1. start session
s_curr,supp_curr = env.reset()
print('Init state:', s_curr.size(), supp_curr.size())

# 2. choose action with some algorithm
chosen_action = torch.tensor([0]*params['BATCH']).to(params['DEVICE']) # 0 = GO UP
print('Actions:', chosen_action.size())

# 3. get next state with enviroment
reward_curr = env.get_reward(s_curr)
s_next, supp_next, reward = env.step(s_curr, supp_curr, chosen_action)
print('Next state:', s_next.size(), supp_next.size(), reward.size(), reward_curr.size())

# если агент появился на камне то вывести:
for k in range(params['BATCH']):
    s0 = (s_curr[k].permute(1,2,0).detach().cpu()+1)/2
    s1 = (s_next[k].permute(1,2,0).detach().cpu()+1)/2
    r_curr = reward_curr[k].detach().cpu()
    r = reward[k].detach().cpu()
    if r_curr==1: # and r==1:
        print('Reward:', r)
        see_plot(torch.cat([s0,s1], dim=1), size=(8,4))<jupyter_output>Init state: torch.Size([128, 3, 96, 96]) torch.Size([128, 2])
Actions: torch.Size([128])
Next state: torch.Size([128, 3, 96, 96]) torch.Size([128, 2]) torch.Size([128, 1]) torch.Size([128, 1])<jupyter_text>DQN model<jupyter_code>class DQN_NET(nn.Module):
    def __init__(self):
        super().__init__()
        self.backbone = models.vgg19(pretrained=True).features
        self.pool = nn.AdaptiveMaxPool2d((1,1), return_indices=False)
        self.head = nn.Sequential(nn.Conv2d(512, 32, kernel_size=(1, 1)),
                                  nn.ReLU(inplace=True),
                                  nn.Dropout2d(0.1),
                                  nn.Conv2d(32, 8, kernel_size=(1, 1)),
                                  )

    def forward(self, x):
        x = self.backbone(x)
        x = self.pool(x)
        x = self.head(x)   
        x = torch.flatten(x, 1)
        return x
r_ts = []
l_ts = []

def dqn_training(model, params):
    model.train()
    MSE = nn.MSELoss()
    opt = optim.Adam(model.parameters(), lr=1e-4)
    env = NeuralEnv("./model_weights/resunet_v5.pth",
                    "./model_weights/nostone_stone_classifier.pth",
                    params['DEVICE'],
                    params['BATCH'],
                    0.95)
    n_attempts = 50
    n_steps = 105
    gamma = 0.7
    alpha = 0.7
    initial_epsilon = 0.2
    final_epsilon = 0.0001
    epsilon_decrements = np.linspace(initial_epsilon, final_epsilon,n_attempts*n_steps)

    for i in range(n_attempts):
        s_curr,supp_curr = env.reset()
        total_reward = 0
        total_loss = 0.0
        for j in range(n_steps):
            iteration = j + i*n_steps
            epsilon = epsilon_decrements[iteration]
            #print(iteration, epsilon)

            value_s_curr = model(s_curr)
            policy_action = torch.argmax(value_s_curr, dim=1) # (batch,)
            random_action = torch.randint(0,8,(params['BATCH'],)).to(params['DEVICE']) # (batch,)
            if random.random()<=epsilon:
                action = random_action
            else:
                action = policy_action
            s_next, supp_next, reward = env.step(s_curr, supp_curr, action)   
            value_s_next = model(s_next)
            y_pred = torch.sum(value_s_curr*F.one_hot(action, num_classes=8).float(), dim=1)
            y_gt = (1-alpha)*y_pred + alpha*(reward.squeeze()+gamma*torch.max(value_s_next, dim=1).values)
            y_gt = y_gt.clone().to(torch.float32).detach()

            opt.zero_grad()
            loss = MSE(y_pred, y_gt) 
            loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0)
            opt.step()

            s_curr = s_next
            supp_curr = supp_next
            total_reward += reward.sum().item()
            total_loss += loss.item()
        
        r_ts.append(total_reward)
        l_ts.append(round(total_loss,4))
        if i % 10 == 0:
            print(i, total_reward, round(total_loss,4))
params = {'DEVICE': 'cuda:0',
          'BATCH': 32,
         'reward_confidence': 0.95}

q_model = DQN_NET().to(params['DEVICE'])
dqn_training(q_model, params)
torch.save(q_model.state_dict(), 'dqn_v1.pth')<jupyter_output>0 21.0 0.4932
10 22.0 0.6168
20 30.0 0.4796
30 37.0 0.5035
40 52.0 0.6239<jupyter_text>Testing<jupyter_code>params = {'DEVICE': 'cuda:0',
          'BATCH': 1,
         'reward_confidence': 0.95}

env = NeuralEnv("./model_weights/resunet_v5.pth",
                "./model_weights/nostone_stone_classifier.pth",
                params['DEVICE'],
                params['BATCH'],
                params['reward_confidence'])

s_init,supp_init = env.reset()
s_curr = s_init[0]
supp_curr = supp_init[0]

see_plot((s_curr.permute(1,2,0).detach().cpu()+1)/2)
inference_agent(q_model, torch.rot90(s_curr, 0, [1,2]), supp_curr) 
def DQN_train():
    # s0,s1,a,r
    #s0 - (batch,3,96,96), (batch,2)
    #s1 - (batch,3,96,96), (batch,2)
    #a -  (batch,8)
    #r -  (batch,)
    
    output0 = model(s0)
    output = model(s1)
    
    y_pred = torch.sum(output0*a, dim=1)
    y_true = (1-alpha)*y_pred + alpha*(r+gamma*torch.max(output, dim=1).values)
    y_true = y_true.clone().to(torch.float32)
    y_true = y_true.detach()
    
    optimizer.zero_grad()
    loss = MSE(y_pred, y_true) 
    loss.backward()
    optimizer.step()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.16161375029840058, 'fraction_numerical': 0.03230683536245723, 'mean_word_length': 2.029893924783028, 'pattern_counts': {'":': 0, '<': 16, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 5, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '177752', 'license': 'MIT', 'n_tokens_mistral': 4473, 'n_tokens_neox': 4260, 'n_words': 769, 'path': 'jupyter/rl/surviv_dqn_v1.ipynb', 'repo_name': 'Laggg/neural_env_surviv', 'size': '12567'}