Model/fur_rl/models/retriever.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from Model.CLIP.cn_clip.clip import load_from_name
import cn_clip.clip as clip
from PIL import Image
from torch.autograd import Variable
from fur_rl.models.StageOne import Stage1
import copy
import PIL

class Retriever(nn.Module):
    def __init__(self, sentence_clip_model2, sentence_clip_preprocess2, device):
        super(Retriever, self).__init__()
        self.sentence_clip_model2 = sentence_clip_model2
        self.sentence_clip_preprocess2 = sentence_clip_preprocess2
        self.tokenize = clip.tokenize
        self.Stage1 = Stage1(hid_dim=512)
        self.conv1 = nn.Conv1d(in_channels=52, out_channels=64, kernel_size=5, stride=1, padding=2, dilation=1, groups=1, bias=True, padding_mode='zeros')
        self.conv2 = nn.Conv1d(in_channels=64, out_channels=32, kernel_size=3, stride=1, padding=2, dilation=1, groups=1, bias=True, padding_mode='zeros')
        self.conv3 = nn.Conv1d(in_channels=32, out_channels=16, kernel_size=3, stride=1, padding=2, dilation=1, groups=1, bias=True, padding_mode='zeros')
        self.conv4 = nn.Conv1d(in_channels=16, out_channels=1, kernel_size=1, stride=1, padding=0, dilation=1, groups=1, bias=True, padding_mode='zeros')
        self.fc1 = nn.Linear(in_features=516, out_features=50, bias=True)
        self.relu = nn.ReLU()
        self.softmax = nn.Softmax(dim=1)
        self.device = device
    def img_embed(self, img):
        img_embed = self.sentence_clip_model2.encode_image(img)
        return img_embed

    def txt_embed(self, txt_token):
        txt_embed = self.sentence_clip_model2.encode_text(txt_token)
        return txt_embed

    def get_env_feature(self, img, txt):
        hx = self.Stage1.hx
        x = self.Stage1(img, txt)
        return x, hx

    def init_hid(self, batch_size=1):
        self.Stage1.init_hid(batch_size=batch_size)
        return

    def detach_hid(self):
        self.Stage1.detach_hid()
        return

    def get_Q(self, img, txt, action):
        env_feature, hx = self.get_env_feature(img, txt)
        env_feature = env_feature.unsqueeze(dim=1)
        hx = hx.unsqueeze(dim=1)
        # cat env_feature and action
        x = torch.cat((env_feature, action, hx), dim=1)
        # convolution layer
        x = self.relu(self.conv1(x))
        x = self.relu(self.conv2(x))
        x = self.relu(self.conv3(x))
        x = self.relu(self.conv4(x))
        # fully connected layer
        x = x.view(x.shape[0], -1)
        x = self.fc1(x)
        # softmax
        # x = self.softmax(x)
        return x

    def forward(self, img, txt, action, hx):
        self.Stage1.hx = hx
        Q = self.get_Q(img, txt, action)
        return Q

# v2 是训练encode的大模型，训练效率太低
class Retriever_v2(nn.Module):
    def __init__(self, sentence_clip_model2, sentence_clip_preprocess2, device):
        super(Retriever_v2, self).__init__()
        self.sentence_clip_model2 = sentence_clip_model2
        self.sentence_clip_preprocess2 = sentence_clip_preprocess2
        self.tokenize = clip.tokenize
        self.Stage1 = Stage1(hid_dim=512)
        self.conv1 = nn.Conv1d(in_channels=22, out_channels=64, kernel_size=5, stride=1, padding=2, dilation=1, groups=1, bias=True, padding_mode='zeros')
        self.conv2 = nn.Conv1d(in_channels=64, out_channels=32, kernel_size=3, stride=1, padding=2, dilation=1, groups=1, bias=True, padding_mode='zeros')
        self.conv3 = nn.Conv1d(in_channels=32, out_channels=16, kernel_size=3, stride=1, padding=2, dilation=1, groups=1, bias=True, padding_mode='zeros')
        self.conv4 = nn.Conv1d(in_channels=16, out_channels=1, kernel_size=1, stride=1, padding=0, dilation=1, groups=1, bias=True, padding_mode='zeros')
        self.fc1 = nn.Linear(in_features=516, out_features=20, bias=True)
        self.relu = nn.ReLU()
        self.softmax = nn.Softmax(dim=1)
        self.device = device
    def img_embed(self, img):
        img_embed = self.sentence_clip_model2.encode_image(img)
        return img_embed

    def txt_embed(self, txt_token):
        txt_embed = self.sentence_clip_model2.encode_text(txt_token)
        return txt_embed

    def get_env_feature(self, img, txt):
        hx = self.Stage1.hx
        x = self.Stage1(img, txt)
        return x, hx

    def init_hid(self, batch_size=1):
        self.Stage1.init_hid(batch_size=batch_size)
        return

    def detach_hid(self):
        self.Stage1.detach_hid()
        return

    def get_Q(self, img, txt, action):
        env_feature, hx = self.get_env_feature(img, txt)
        env_feature = env_feature.unsqueeze(dim=1)
        hx = hx.unsqueeze(dim=1)
        # cat env_feature and action
        x = torch.cat((env_feature, action, hx), dim=1)
        # convolution layer
        x = self.relu(self.conv1(x))
        x = self.relu(self.conv2(x))
        x = self.relu(self.conv3(x))
        x = self.relu(self.conv4(x))
        # fully connected layer
        x = x.view(x.shape[0], -1)
        x = self.fc1(x)
        # softmax
        # x = self.softmax(x)
        return x

    def forward(self, img, txt, action, hx):
        # img txt 是最原始的图片和文本
        self.Stage1.hx = hx
        img_embed = self.img_embed(img)     # 2 512
        txt_embed = self.txt_embed(self.tokenize(txt).to(self.device))
        # print("txt_embed", img_embed.shape, txt_embed.shape)
        c = []
        for i in range(action.shape[1]):
            action_i = self.img_embed(action[:, i])
            c.append(action_i)
        c = torch.stack(c, dim=1)   # 2 20 512
        Q = self.get_Q(img_embed, txt_embed, c)
        return Q
# v3 带softmax的模型，actor
class Retriever_v3(nn.Module):
    def __init__(self, sentence_clip_model2, sentence_clip_preprocess2, device):
        super(Retriever_v3, self).__init__()
        self.sentence_clip_model2 = sentence_clip_model2
        self.sentence_clip_preprocess2 = sentence_clip_preprocess2
        self.tokenize = clip.tokenize
        self.Stage1 = Stage1(hid_dim=512)
        self.conv1 = nn.Conv1d(in_channels=52, out_channels=64, kernel_size=5, stride=1, padding=2, dilation=1, groups=1, bias=True, padding_mode='zeros')
        self.conv2 = nn.Conv1d(in_channels=64, out_channels=32, kernel_size=3, stride=1, padding=2, dilation=1, groups=1, bias=True, padding_mode='zeros')
        self.conv3 = nn.Conv1d(in_channels=32, out_channels=16, kernel_size=3, stride=1, padding=2, dilation=1, groups=1, bias=True, padding_mode='zeros')
        self.conv4 = nn.Conv1d(in_channels=16, out_channels=1, kernel_size=1, stride=1, padding=0, dilation=1, groups=1, bias=True, padding_mode='zeros')
        self.bn1 = nn.BatchNorm1d(64)
        self.bn2 = nn.BatchNorm1d(32)
        self.bn3 = nn.BatchNorm1d(16)
        self.bn4 = nn.BatchNorm1d(1)
        self.fc1 = nn.Linear(in_features=516, out_features=50, bias=True)
        self.bn5 = nn.BatchNorm1d(50)
        # self.fc2 = nn.Linear(in_features=516, out_features=50, bias=True)
        # self.fc3 = nn.Linear(in_features=516, out_features=10, bias=True)

        self.relu = nn.ReLU()
        self.softmax = nn.Softmax(dim=1)
        self.device = device
    def img_embed(self, img):
        img_embed = self.sentence_clip_model2.encode_image(img)
        img_embed = img_embed / img_embed.norm(dim=-1, keepdim=True)
        return img_embed

    def txt_embed(self, txt_token):
        txt_embed = self.sentence_clip_model2.encode_text(txt_token)
        txt_embed = txt_embed / txt_embed.norm(dim=-1, keepdim=True)
        return txt_embed

    def get_env_feature(self, img, txt, reward):
        hx = self.Stage1.hx
        x = self.Stage1(img, txt, reward)
        return x, hx

    def init_hid(self, batch_size=1):
        self.Stage1.init_hid(batch_size=batch_size)
        return

    def detach_hid(self):
        self.Stage1.detach_hid()
        return

    def get_Q(self, img, txt, action, reward):
        env_feature, hx = self.get_env_feature(img, txt, reward)
        env_feature = env_feature.unsqueeze(dim=1)
        hx = hx.unsqueeze(dim=1)
        # cat env_feature and action
        x = torch.cat((env_feature, action, hx), dim=1)
        # convolution layer
        x = self.relu(self.bn1(self.conv1(x)))
        x = self.relu(self.bn2(self.conv2(x)))
        x = self.relu(self.bn3(self.conv3(x)))
        x = self.relu(self.bn4(self.conv4(x)))
        # fully connected layer
        x = x.view(x.shape[0], -1)
        x = self.fc1(x)
        x = self.bn5(x)
        # softmax
        x = self.softmax(x)
        return x

    def forward(self, img, txt, action, hx, r):
        self.Stage1.hx = hx
        Q = self.get_Q(img, txt, action, r)
        return Q

# v4 删除了hx输入后面卷积，测试hx对记忆的影响
class Retriever_v4(nn.Module):
    def __init__(self, sentence_clip_model2, sentence_clip_preprocess2, device):
        super(Retriever_v4, self).__init__()
        self.sentence_clip_model2 = sentence_clip_model2
        self.sentence_clip_preprocess2 = sentence_clip_preprocess2
        self.tokenize = clip.tokenize
        self.Stage1 = Stage1(hid_dim=512)
        self.conv1 = nn.Conv1d(in_channels=51, out_channels=64, kernel_size=5, stride=1, padding=2, dilation=1, groups=1, bias=True, padding_mode='zeros')
        self.conv2 = nn.Conv1d(in_channels=64, out_channels=32, kernel_size=3, stride=1, padding=2, dilation=1, groups=1, bias=True, padding_mode='zeros')
        self.conv3 = nn.Conv1d(in_channels=32, out_channels=16, kernel_size=3, stride=1, padding=2, dilation=1, groups=1, bias=True, padding_mode='zeros')
        self.conv4 = nn.Conv1d(in_channels=16, out_channels=1, kernel_size=1, stride=1, padding=0, dilation=1, groups=1, bias=True, padding_mode='zeros')
        self.fc1 = nn.Linear(in_features=516, out_features=50, bias=True)
        self.relu = nn.ReLU()
        self.softmax = nn.Softmax(dim=1)
        self.device = device
    def img_embed(self, img):
        img_embed = self.sentence_clip_model2.encode_image(img)
        return img_embed

    def txt_embed(self, txt_token):
        txt_embed = self.sentence_clip_model2.encode_text(txt_token)
        return txt_embed

    def get_env_feature(self, img, txt):
        hx = self.Stage1.hx
        x = self.Stage1(img, txt)
        return x, hx

    def init_hid(self, batch_size=1):
        self.Stage1.init_hid(batch_size=batch_size)
        return

    def detach_hid(self):
        self.Stage1.detach_hid()
        return

    def get_Q(self, img, txt, action):
        env_feature, hx = self.get_env_feature(img, txt)
        env_feature = env_feature.unsqueeze(dim=1)
        hx = hx.unsqueeze(dim=1)
        # cat env_feature and action
        x = torch.cat((env_feature, action), dim=1)
        # convolution layer
        x = self.relu(self.conv1(x))
        x = self.relu(self.conv2(x))
        x = self.relu(self.conv3(x))
        x = self.relu(self.conv4(x))
        # fully connected layer
        x = x.view(x.shape[0], -1)
        x = self.fc1(x)
        # softmax
        x = self.softmax(x)
        return x

    def forward(self, img, txt, action, hx):
        self.Stage1.hx = hx
        Q = self.get_Q(img, txt, action)
        return Q
# v5 critic模型，只用1个candidate
class Retriever_v5(nn.Module):
    def __init__(self, sentence_clip_model2, sentence_clip_preprocess2, device):
        super(Retriever_v5, self).__init__()
        self.sentence_clip_model2 = sentence_clip_model2
        self.sentence_clip_preprocess2 = sentence_clip_preprocess2
        self.tokenize = clip.tokenize
        self.Stage1 = Stage1(hid_dim=512)
        # self.conv1 = nn.Conv1d(in_channels=3, out_channels=1, kernel_size=3, stride=1, padding=2, dilation=1, groups=1, bias=True, padding_mode='zeros')
        # self.conv2 = nn.Conv1d(in_channels=64, out_channels=32, kernel_size=3, stride=1, padding=2, dilation=1, groups=1, bias=True, padding_mode='zeros')
        # self.conv3 = nn.Conv1d(in_channels=32, out_channels=16, kernel_size=3, stride=1, padding=2, dilation=1, groups=1, bias=True, padding_mode='zeros')
        # self.conv4 = nn.Conv1d(in_channels=16, out_channels=1, kernel_size=1, stride=1, padding=0, dilation=1, groups=1, bias=True, padding_mode='zeros')
        self.fc01 = nn.Linear(in_features=512 * 2, out_features=512, bias=True)
        self.fc02 = nn.Linear(in_features=512, out_features=256, bias=True)
        self.fc03 = nn.Linear(in_features=256, out_features=64, bias=True)
        self.fc04 = nn.Linear(in_features=64, out_features=1, bias=True)
        # self.fc1 = nn.Linear(in_features=516, out_features=1, bias=True)
        self.fc2 = nn.Linear(in_features=512, out_features=512, bias=True)
        self.relu = nn.ReLU()
        self.softmax = nn.Softmax(dim=1)
        self.device = device
    def img_embed(self, img):
        img_embed = self.sentence_clip_model2.encode_image(img)
        img_embed = img_embed / img_embed.norm(dim=-1, keepdim=True)
        return img_embed

    def txt_embed(self, txt_token):
        txt_embed = self.sentence_clip_model2.encode_text(txt_token)
        txt_embed = txt_embed / txt_embed.norm(dim=-1, keepdim=True)
        return txt_embed

    def get_env_feature(self, img, txt, reward):
        hx = self.Stage1.hx
        x = self.Stage1(img, txt, reward)
        x = x / x.norm(dim=-1, keepdim=True)
        return x, hx

    def init_hid(self, batch_size=1):
        self.Stage1.init_hid(batch_size=batch_size)
        return

    def detach_hid(self):
        self.Stage1.detach_hid()
        return

    def get_Q(self, img, txt, action, reward):
        env_feature, hx = self.get_env_feature(img, txt, reward)
        env_feature = env_feature.unsqueeze(dim=1)
        hx = hx.unsqueeze(dim=1)
        # cat env_feature and action
        action = action.squeeze(dim=1)
        action = self.fc2(action)
        action = action.unsqueeze(dim=1)
        # x = torch.cat((env_feature, action, hx), dim=1)
        # print(env_feature.shape, action.shape)
        x = torch.cat((env_feature, action), dim=2)
        # convolution layer
        # print(x.shape)
        x = self.fc01(x)
        # x = self.relu(self.conv1(x))
        # print(x.shape)
        x = self.fc02(x)
        x = self.fc03(x)
        x = self.fc04(x)
        x = x.squeeze(dim=1)
        # print(x.shape)
        # x = self.relu(self.conv1(x))
        # x = self.relu(self.conv2(x))
        # x = self.relu(self.conv3(x))
        # x = self.relu(self.conv4(x))
        # fully connected layer
        # x = x.view(x.shape[0], -1)
        # x = self.fc1(x)
        return x

    def forward(self, img, txt, action, hx, r):
        self.Stage1.hx = hx
        Q = self.get_Q(img, txt, action, r)
        return Q
# v6 训练一个选择candidate的模型
class Retriever_v6(nn.Module):
    def __init__(self, sentence_clip_model2, sentence_clip_preprocess2, device):
        super(Retriever_v6, self).__init__()
        self.sentence_clip_model2 = sentence_clip_model2
        self.sentence_clip_preprocess2 = sentence_clip_preprocess2
        self.tokenize = clip.tokenize
        self.Stage1 = Stage1(hid_dim=512)
        self.conv1 = nn.Conv1d(in_channels=52, out_channels=64, kernel_size=5, stride=1, padding=2, dilation=1, groups=1, bias=True, padding_mode='zeros')
        self.conv2 = nn.Conv1d(in_channels=64, out_channels=32, kernel_size=3, stride=1, padding=2, dilation=1, groups=1, bias=True, padding_mode='zeros')
        self.conv3 = nn.Conv1d(in_channels=32, out_channels=16, kernel_size=3, stride=1, padding=2, dilation=1, groups=1, bias=True, padding_mode='zeros')
        self.conv4 = nn.Conv1d(in_channels=16, out_channels=1, kernel_size=1, stride=1, padding=0, dilation=1, groups=1, bias=True, padding_mode='zeros')
        self.fc1 = nn.Linear(in_features=516, out_features=50, bias=True)
        self.relu = nn.ReLU()
        self.softmax = nn.Softmax(dim=1)
        self.device = device
    def img_embed(self, img):
        img_embed = self.sentence_clip_model2.encode_image(img)
        return img_embed

    def txt_embed(self, txt_token):
        txt_embed = self.sentence_clip_model2.encode_text(txt_token)
        return txt_embed

    def get_env_feature(self, img, txt):
        hx = self.Stage1.hx
        x = self.Stage1(img, txt)
        return x, hx

    def init_hid(self, batch_size=1):
        self.Stage1.init_hid(batch_size=batch_size)
        return

    def detach_hid(self):
        self.Stage1.detach_hid()
        return

    def get_Q(self, img, txt, action):
        env_feature, hx = self.get_env_feature(img, txt)
        env_feature = env_feature.unsqueeze(dim=1)
        hx = hx.unsqueeze(dim=1)
        # cat env_feature and action
        x = torch.cat((env_feature, action, hx), dim=1)
        # convolution layer
        x = self.relu(self.conv1(x))
        x = self.relu(self.conv2(x))
        x = self.relu(self.conv3(x))
        x = self.relu(self.conv4(x))
        # fully connected layer
        x = x.view(x.shape[0], -1)
        x = self.fc1(x)
        # softmax
        x = self.softmax(x)
        return x

    def forward(self, img, txt, action, hx):
        self.Stage1.hx = hx
        Q = self.get_Q(img, txt, action)
        return Q



MEMORY_CAPACITY = 100
N_STATES = 4
GAMMA = 0.95

class DQN():
    def __init__(self, sentence_clip_model2, sentence_clip_preprocess2, device):
        self.eval_net = Retriever_v5(sentence_clip_model2, sentence_clip_preprocess2, device).to(device)
        self.target_net = Retriever_v5(sentence_clip_model2, sentence_clip_preprocess2, device).to(device)
        self.num_train = 0

        self.memory = []     # initialize memory
        self.memory_counter = 0
        self.optimizer = torch.optim.Adam(self.eval_net.parameters(), lr=0.01)
        self.loss_func = nn.MSELoss()

    def store_transition(self, g, f, c, hx, a, r, g_, f_, c_, hx_, batch_size):
        for i in range(batch_size):
            if r[i] == -10000:
                continue
            if r[i] > 100:
                r[i] = 100
            self.memory.append((g[i], f[i], c[i], hx[i], a[i], r[i], g_[i], f_[i], c_[i], hx_[i]))
            if len(self.memory) > MEMORY_CAPACITY:
                self.memory.pop(0)
            self.memory_counter += 1

    def learn(self, batch_size, device=None):
        # target parameter update
        if self.memory_counter > MEMORY_CAPACITY:
            sample_index = np.random.choice(MEMORY_CAPACITY, batch_size)
        else:
            sample_index = np.random.choice(self.memory_counter, batch_size)
        # print("sample_index", sample_index)

        b_g, b_f, b_c, b_hx, b_a, b_r, b_g_, b_f_, b_c_, b_hx_ = self.memory[sample_index[0]]
        b_g = b_g.unsqueeze(dim=0)
        b_f = b_f.unsqueeze(dim=0)
        b_c = b_c.unsqueeze(dim=0)
        b_hx = b_hx.unsqueeze(dim=0)
        b_a = b_a.unsqueeze(dim=0)
        b_r = b_r.unsqueeze(dim=0)
        b_g_ = b_g_.unsqueeze(dim=0)
        b_f_ = b_f_.unsqueeze(dim=0)
        b_c_ = b_c_.unsqueeze(dim=0)
        b_hx_ = b_hx_.unsqueeze(dim=0)
        for i in range(1, len(sample_index)):
            g, f, c, hx, a, r, g_, f_, c_, hx_ = self.memory[sample_index[i]]
            b_g = torch.cat((b_g, g.unsqueeze(dim=0)), dim=0)
            b_f = torch.cat((b_f, f.unsqueeze(dim=0)), dim=0)
            b_c = torch.cat((b_c, c.unsqueeze(dim=0)), dim=0)
            b_hx = torch.cat((b_hx, hx.unsqueeze(dim=0)), dim=0)
            b_a = torch.cat((b_a, a.unsqueeze(dim=0)), dim=0)
            b_r = torch.cat((b_r, r.unsqueeze(dim=0)), dim=0)
            b_g_ = torch.cat((b_g_, g_.unsqueeze(dim=0)), dim=0)
            b_f_ = torch.cat((b_f_, f_.unsqueeze(dim=0)), dim=0)
            b_c_ = torch.cat((b_c_, c_.unsqueeze(dim=0)), dim=0)
            b_hx_ = torch.cat((b_hx_, hx_.unsqueeze(dim=0)), dim=0)
        # print("bhx.shape", b_hx.shape)
        b_a = b_a.unsqueeze(dim=1)
        b_r = b_r.unsqueeze(dim=1)
        #
        # if self.num_train > 5000:
        #     self.target_net.load_state_dict(self.eval_net.state_dict())
        #     # print("target_net update...")
        #     self.num_train = 0
        # self.num_train += 1
        # print("shape", b_g.shape, b_f.shape, b_c.shape, b_hx.shape, b_a.shape, b_r.shape, b_g_.shape, b_f_.shape, b_c_.shape, b_hx_.shape)
        # # DQN multi
        # q_eval = self.eval_net(b_g, b_f.detach(), b_c, b_hx.detach()).gather(1, b_a)  # shape (batch, 1)
        # q_next = self.target_net(b_g_, b_f_, b_c_, b_hx_).detach()     # detach from graph, don't backpropagate
        # q_target = b_r + GAMMA * q_next.max(1)[0].view(batch_size, 1)   # shape (batch, 1)

        # DQN single, retrieval-v5
        b_c1 = torch.zeros(b_c.shape[0], 1, b_c.shape[2]).to(device)
        for i in range(b_a.shape[0]):
            b_c1[i][0] = b_c[i][b_a[i][0]].to(device)
        q_eval = self.eval_net(b_g, b_f.detach(), b_c1, b_hx.detach())  # shape (batch, 1)
        q_next = torch.zeros((batch_size, len(b_c_[0]))).to(device)
        for i in range(len(b_c_[0])):
            q_temp = self.target_net(b_g_, b_f_, b_c_[:, i, :].unsqueeze(1), b_hx_).detach()
            q_next[:, i] = q_temp.squeeze(dim=1)
        flag = 0
        for j in range(len(b_c_)):  # batch_i
            if b_r[j] >= 95:
                q_next[j, :] = torch.zeros((1,len(b_c_[0]))).to(device)
                flag = 1
        q_target = b_r + GAMMA * q_next.max(1)[0].view(batch_size, 1)  # shape (batch, 1)
        # print("q_target", q_eval, q_target, b_r)
        loss = self.loss_func(q_eval.float(), q_target.float())
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()
        return loss


class DQN_v1():
    def __init__(self, sentence_clip_model2, sentence_clip_preprocess2, device):
        self.eval_net = Retriever_v5(sentence_clip_model2, sentence_clip_preprocess2, device).to(device)
        self.target_net = Retriever_v5(sentence_clip_model2, sentence_clip_preprocess2, device).to(device)
        self.num_train = 0

        self.memory = []     # initialize memory
        self.memory_counter = 0
        self.optimizer = torch.optim.Adam(self.eval_net.parameters(), lr=0.01)
        self.loss_func = nn.MSELoss()

    def store_transition(self, g, f, c, hx, a, r, g_, f_, c_, hx_, batch_size):
        for i in range(batch_size):
            if r[i] == -10000:
                continue
            self.memory.append((g[i], f[i], c[i], hx[i], a[i], r[i], g_[i], f_[i], c_[i], hx_[i]))
            if len(self.memory) > MEMORY_CAPACITY:
                self.memory.pop(0)
            self.memory_counter += 1

    def learn(self, batch_size, device=None):
        # target parameter update
        if self.memory_counter > MEMORY_CAPACITY:
            sample_index = np.random.choice(MEMORY_CAPACITY, batch_size)
        else:
            sample_index = np.random.choice(self.memory_counter, batch_size)
        # print("sample_index", sample_index)

        b_g, b_f, b_c, b_hx, b_a, b_r, b_g_, b_f_, b_c_, b_hx_ = self.memory[sample_index[0]]
        b_g = b_g.unsqueeze(dim=0)
        b_fl = []
        b_fl.append(b_f)
        # b_f = b_f.unsqueeze(dim=0)
        b_c = b_c.unsqueeze(dim=0)
        b_hx = b_hx.unsqueeze(dim=0)
        b_a = b_a.unsqueeze(dim=0)
        b_r = b_r.unsqueeze(dim=0)
        b_g_ = b_g_.unsqueeze(dim=0)
        b_fl_ = []
        b_fl_.append(b_f_)
        # b_f_ = b_f_.unsqueeze(dim=0)
        b_c_ = b_c_.unsqueeze(dim=0)
        b_hx_ = b_hx_.unsqueeze(dim=0)
        for i in range(1, len(sample_index)):
            g, f, c, hx, a, r, g_, f_, c_, hx_ = self.memory[sample_index[i]]
            b_g = torch.cat((b_g, g.unsqueeze(dim=0)), dim=0)
            b_fl.append(f)
            # b_f = torch.cat((b_f, f.unsqueeze(dim=0)), dim=0)
            b_c = torch.cat((b_c, c.unsqueeze(dim=0)), dim=0)
            b_hx = torch.cat((b_hx, hx.unsqueeze(dim=0)), dim=0)
            b_a = torch.cat((b_a, a.unsqueeze(dim=0)), dim=0)
            b_r = torch.cat((b_r, r.unsqueeze(dim=0)), dim=0)
            b_g_ = torch.cat((b_g_, g_.unsqueeze(dim=0)), dim=0)
            b_fl_.append(f_)
            # b_f_ = torch.cat((b_f_, f_.unsqueeze(dim=0)), dim=0)
            b_c_ = torch.cat((b_c_, c_.unsqueeze(dim=0)), dim=0)
            b_hx_ = torch.cat((b_hx_, hx_.unsqueeze(dim=0)), dim=0)
        # print("bhx.shape", b_hx.shape)
        b_a = b_a.unsqueeze(dim=1)
        b_r = b_r.unsqueeze(dim=1)

        if self.num_train > 5000:
            self.target_net.load_state_dict(self.eval_net.state_dict())
            self.num_train = 0
        self.num_train += 1
        # DQN multi,
        q_eval = self.eval_net(b_g, b_fl, b_c, b_hx.detach()).gather(1, b_a)  # shape (batch, 1)
        with torch.no_grad():
            q_next = self.target_net(b_g_, b_fl_, b_c_, b_hx_).detach()
        q_next = self.target_net(b_g_, b_fl_, b_c_, b_hx_).detach()     # detach from graph, don't backpropagate
        q_target = b_r + GAMMA * q_next.max(1)[0].view(batch_size, 1)   # shape (batch, 1)

        loss = self.loss_func(q_eval.float(), q_target.float())
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()
        return loss

# 使用actor critic模型
class DQN_v2():
    def __init__(self, sentence_clip_model2, sentence_clip_preprocess2, device):
        self.eval_net = Retriever(sentence_clip_model2, sentence_clip_preprocess2, device).to(device)
        self.target_net = Retriever(sentence_clip_model2, sentence_clip_preprocess2, device).to(device)
        self.actor_net = Retriever_v3(sentence_clip_model2, sentence_clip_preprocess2, device).to(device)
        self.num_train = 0

        self.memory = []     # initialize memory
        self.memory_counter = 0
        self.actor_optimizer = torch.optim.Adam(self.actor_net.parameters(), lr=0.001)
        self.critic_optimizer = torch.optim.Adam(self.eval_net.parameters(), lr=0.001)
        self.loss_func = nn.MSELoss()

    def store_transition(self, g, f, c, hx, a, r, g_, f_, c_, hx_, batch_size):
        for i in range(batch_size):
            if r[i] == -10000:
                continue
            self.memory.append((g[i], f[i], c[i], hx[i], a[i], r[i], g_[i], f_[i], c_[i], hx_[i]))
            if len(self.memory) > MEMORY_CAPACITY:
                self.memory.pop(0)
            self.memory_counter += 1

    def learn(self, batch_size):
        # target parameter update
        if self.memory_counter > MEMORY_CAPACITY:
            sample_index = np.random.choice(MEMORY_CAPACITY, batch_size)
        else:
            sample_index = np.random.choice(self.memory_counter, batch_size)
        # print("sample_index", sample_index)

        b_g, b_f, b_c, b_hx, b_a, b_r, b_g_, b_f_, b_c_, b_hx_ = self.memory[sample_index[0]]
        b_g = b_g.unsqueeze(dim=0)
        b_f = b_f.unsqueeze(dim=0)
        b_c = b_c.unsqueeze(dim=0)
        b_hx = b_hx.unsqueeze(dim=0)
        b_a = b_a.unsqueeze(dim=0)
        b_r = b_r.unsqueeze(dim=0)
        b_g_ = b_g_.unsqueeze(dim=0)
        b_f_ = b_f_.unsqueeze(dim=0)
        b_c_ = b_c_.unsqueeze(dim=0)
        b_hx_ = b_hx_.unsqueeze(dim=0)
        for i in range(1, len(sample_index)):
            g, f, c, hx, a, r, g_, f_, c_, hx_ = self.memory[sample_index[i]]
            b_g = torch.cat((b_g, g.unsqueeze(dim=0)), dim=0)
            b_f = torch.cat((b_f, f.unsqueeze(dim=0)), dim=0)
            b_c = torch.cat((b_c, c.unsqueeze(dim=0)), dim=0)
            b_hx = torch.cat((b_hx, hx.unsqueeze(dim=0)), dim=0)
            b_a = torch.cat((b_a, a.unsqueeze(dim=0)), dim=0)
            b_r = torch.cat((b_r, r.unsqueeze(dim=0)), dim=0)
            b_g_ = torch.cat((b_g_, g_.unsqueeze(dim=0)), dim=0)
            b_f_ = torch.cat((b_f_, f_.unsqueeze(dim=0)), dim=0)
            b_c_ = torch.cat((b_c_, c_.unsqueeze(dim=0)), dim=0)
            b_hx_ = torch.cat((b_hx_, hx_.unsqueeze(dim=0)), dim=0)
        # print("bhx.shape", b_hx.shape)
        b_a = b_a.unsqueeze(dim=1)
        b_r = b_r.unsqueeze(dim=1)

        if self.num_train > 50:
            self.target_net.load_state_dict(self.eval_net.state_dict())
            self.num_train = 0
        self.num_train += 1
        # print("shape", b_g.shape, b_f.shape, b_c.shape, b_hx.shape, b_a.shape, b_r.shape, b_g_.shape, b_f_.shape, b_c_.shape, b_hx_.shape)
        q_eval = self.eval_net(b_g, b_f.detach(), b_c, b_hx.detach()).gather(1, b_a)  # shape (batch, 1)
        q_next = self.target_net(b_g_, b_f_, b_c_, b_hx_).detach()     # detach from graph, don't backpropagate
        q_target = b_r + GAMMA * q_next.max(1)[0].view(batch_size, 1)   # shape (batch, 1)

        critic_loss = torch.mean(F.mse_loss(q_eval.float(), q_target.float()))

        log_probs = self.actor_net(b_g, b_f.detach(), b_c, b_hx.detach())
        delta = q_target - q_eval
        # print("log_probs", log_probs)
        actor_loss = torch.mean(-log_probs * delta.detach()).float()

        self.actor_optimizer.zero_grad()
        self.critic_optimizer.zero_grad()
        actor_loss.backward()
        critic_loss.backward()
        self.actor_optimizer.step()
        self.critic_optimizer.step()
        # self.target_net.load_state_dict(self.eval_net.state_dict())
        return actor_loss, critic_loss


class DQN_v3():
    def __init__(self, sentence_clip_model2, sentence_clip_preprocess2, device, MULTI_GPU=False, device_ids=None):
        self.eval_net = Retriever_v5(sentence_clip_model2, sentence_clip_preprocess2, device).to(device)
        self.target_net = Retriever_v5(sentence_clip_model2, sentence_clip_preprocess2, device).to(device)
        self.actor_net = Retriever_v3(sentence_clip_model2, sentence_clip_preprocess2, device).to(device)
        self.num_train = 0
        self.device = device
        self.MULTI_GPU = MULTI_GPU
        self.device_ids = device_ids

        self.memory = []  # initialize memory
        self.memory_counter = 0
        self.actor_optimizer = torch.optim.Adam(self.actor_net.parameters(), lr=0.0001)
        self.critic_optimizer = torch.optim.Adam(self.eval_net.parameters(), lr=0.001)
        self.loss_func = nn.MSELoss()


    def store_transition(self, g, f, c, hx, a, r, g_, f_, c_, hx_, batch_size, net_mem):
        for i in range(batch_size):
            if r[i] <= -2000:
                continue
            else:
                g_tmp = copy.deepcopy(g[i])
                f_tmp = copy.deepcopy(f[i])
                c_tmp = copy.deepcopy(c[i])
                hx_tmp = copy.deepcopy(hx[i])
                a_tmp = copy.deepcopy(a[i])
                reward_temp = copy.deepcopy(r[i])
                g_tmp_ = copy.deepcopy(g_[i])
                f_tmp_ = copy.deepcopy(f_[i])
                c_tmp_ = copy.deepcopy(c_[i])
                hx_tmp_ = copy.deepcopy(hx_[i])
                self.memory.append((g_tmp, f_tmp, c_tmp, hx_tmp, a_tmp, reward_temp, g_tmp_, f_tmp_, c_tmp_, hx_tmp_))
                if len(self.memory) > net_mem:
                    self.memory.pop(0)
                self.memory_counter = len(self.memory)
    def learn(self, batch_size, device=None):
        # target parameter update
        if self.memory_counter < batch_size:
            return 0, 0
        else:
            sample_index = np.random.choice(len(self.memory), batch_size, replace=False)
        sample_index = torch.LongTensor(sample_index).to(device)
        g, f, c, hx, a, r, g_, f_, c_, hx_ = zip(*[self.memory[i] for i in sample_index])
        b_g = torch.stack(g)
        b_f = torch.stack(f)
        b_c = torch.stack(c)
        b_hx = torch.stack(hx)
        b_a = torch.stack(a).unsqueeze(1)
        b_r = torch.stack(r)
        b_g_ = torch.stack(g_)
        b_f_ = torch.stack(f_)
        b_c_ = torch.stack(c_)
        b_hx_ = torch.stack(hx_)
        # print("b_g", b_g.shape, b_f.shape, b_c.shape, b_hx.shape, b_a.shape, b_r.shape, b_g_.shape, b_f_.shape, b_c_.shape, b_hx_.shape)
        if self.num_train > 0:
            self.target_net.load_state_dict(self.eval_net.state_dict())
            self.num_train = 0
        self.num_train += 1
        b_c1 = torch.zeros(b_c.shape[0], 1, b_c.shape[2]).to(device)
        for i in range(b_a.shape[0]):
            b_c1[i][0] = b_c[i][b_a[i][0]].to(device)
        q_eval = self.eval_net(b_g, b_f.detach(), b_c1, b_hx.detach(), b_r)  # shape (batch, 1)
        q_next = torch.zeros((batch_size, len(b_c_[0]))).to(device)
        for i in range(len(b_c_[0])):
            q_temp = self.target_net(b_g_, b_f_, b_c_[:, i, :].unsqueeze(1), b_hx_, b_r).detach()
            q_next[:, i] = q_temp.squeeze(dim=1)
        for j in range(len(b_c_)):
            if b_r[j] >= 900:
                q_next[j, :] = torch.zeros((1, len(b_c_[0]))).to(device)
            if b_r[j] <= -900:
                q_next[j, :] = torch.zeros((1, len(b_c_[0]))).to(device)
        q_target = b_r + 0.995 * q_next.max(1)[0].view(batch_size, 1)   # shape (batch, 1)
        print("q_eval, q_target", q_eval[0], q_target[0], b_r[0], q_eval[0] - q_target[0])
        critic_loss = self.loss_func(q_eval.float(), q_target.float())
        log_probs = self.actor_net(b_g, b_f.detach(), b_c, b_hx.detach(), b_r)
        # print("log_probs", log_probs.shape, log_probs)
        log_probs_temp = torch.zeros((batch_size, 1)).to(device)
        for i in range(b_a.shape[0]):
            log_probs_temp[i][0] = log_probs[i][b_a[i][0]].to(device)
            if log_probs_temp[i][0] > 0.9:
                log_probs_temp[i][0] -= 0.1
        # print("log_probs_temp", log_probs_temp.shape, log_probs_temp)

        delta = q_target - q_eval
        for j in range(len(b_c_)):
            if b_r[j] >= 100:
                delta[j] = b_r[j].to(device)
            if b_r[j] <= -900:
                delta[j] = b_r[j].to(device)
        delta = delta.detach()
        b_r = b_r.detach()
        # print("log_probs_temp", log_probs_temp)
        # print("delta", delta)
        # print("b_r", b_r)
        actor_loss = torch.mean( - torch.log(log_probs_temp) * delta).float()
        # print("actor_loss", actor_loss)
        self.actor_optimizer.zero_grad()
        actor_loss.backward()
        self.actor_optimizer.step()

        self.critic_optimizer.zero_grad()
        critic_loss.backward()
        self.critic_optimizer.step()

        return actor_loss, critic_loss


# 打开txt的encode
class A_C_txt():
    def __init__(self, sentence_clip_model2, sentence_clip_preprocess2, device):
        self.eval_net = Retriever_v5(sentence_clip_model2, sentence_clip_preprocess2, device).to(device)
        self.target_net = Retriever_v5(sentence_clip_model2, sentence_clip_preprocess2, device).to(device)
        self.actor_net = Retriever_v3(sentence_clip_model2, sentence_clip_preprocess2, device).to(device)
        self.num_train = 0

        self.memory = []     # initialize memory
        self.memory_counter = 0
        self.actor_optimizer = torch.optim.Adam(self.actor_net.parameters(), lr=0.0001)
        self.critic_optimizer = torch.optim.Adam(self.eval_net.parameters(), lr=0.0001)
        self.loss_func = nn.MSELoss()

    def store_transition(self, g, f, c, hx, a, r, g_, f_, c_, hx_, batch_size, txt_token):
        for i in range(batch_size):
            if r[i] <= -1000:
                continue
            self.memory.append((g[i], f[i], c[i], hx[i], a[i], r[i], g_[i], f_[i], c_[i], hx_[i], txt_token[i]))
            if len(self.memory) > MEMORY_CAPACITY:
                self.memory.pop(0)
            self.memory_counter += 1

    def learn(self, batch_size, device=None):
        # target parameter update
        if self.memory_counter > MEMORY_CAPACITY:
            sample_index = np.random.choice(MEMORY_CAPACITY, batch_size)
        else:
            sample_index = np.random.choice(self.memory_counter, batch_size)
        # print("sample_index", sample_index)

        b_g, b_f, b_c, b_hx, b_a, b_r, b_g_, b_f_, b_c_, b_hx_, txt_token = self.memory[sample_index[0]]
        b_g = b_g.unsqueeze(dim=0)
        b_f = b_f.unsqueeze(dim=0)
        b_c = b_c.unsqueeze(dim=0)
        b_hx = b_hx.unsqueeze(dim=0)
        b_a = b_a.unsqueeze(dim=0)
        b_r = b_r.unsqueeze(dim=0)
        b_g_ = b_g_.unsqueeze(dim=0)
        b_f_ = b_f_.unsqueeze(dim=0)
        b_c_ = b_c_.unsqueeze(dim=0)
        b_hx_ = b_hx_.unsqueeze(dim=0)
        b_txt_token = txt_token.unsqueeze(dim=0)
        for i in range(1, len(sample_index)):
            g, f, c, hx, a, r, g_, f_, c_, hx_ = self.memory[sample_index[i]]
            b_g = torch.cat((b_g, g.unsqueeze(dim=0)), dim=0)
            b_f = torch.cat((b_f, f.unsqueeze(dim=0)), dim=0)
            b_c = torch.cat((b_c, c.unsqueeze(dim=0)), dim=0)
            b_hx = torch.cat((b_hx, hx.unsqueeze(dim=0)), dim=0)
            b_a = torch.cat((b_a, a.unsqueeze(dim=0)), dim=0)
            b_r = torch.cat((b_r, r.unsqueeze(dim=0)), dim=0)
            b_g_ = torch.cat((b_g_, g_.unsqueeze(dim=0)), dim=0)
            b_f_ = torch.cat((b_f_, f_.unsqueeze(dim=0)), dim=0)
            b_c_ = torch.cat((b_c_, c_.unsqueeze(dim=0)), dim=0)
            b_hx_ = torch.cat((b_hx_, hx_.unsqueeze(dim=0)), dim=0)

        # print("bhx.shape", b_hx.shape)
        b_a = b_a.unsqueeze(dim=1)
        # b_r = b_r.unsqueeze(dim=1)

        if self.num_train > 0:
            self.target_net.load_state_dict(self.eval_net.state_dict())
            self.num_train = 0
        self.num_train += 1
        # print("shape", b_g.shape, b_f.shape, b_c.shape, b_hx.shape, b_a.shape, b_r.shape, b_g_.shape, b_f_.shape, b_c_.shape, b_hx_.shape)
        b_c1 = torch.zeros(b_c.shape[0], 1, b_c.shape[2]).to(device)
        for i in range(b_a.shape[0]):
            b_c1[i][0] = b_c[i][b_a[i][0]].to(device)
        q_eval = self.eval_net(b_g, b_f.detach(), b_c1, b_hx.detach(), b_r)  # shape (batch, 1)
        q_next = torch.zeros((batch_size, len(b_c_[0]))).to(device)
        for i in range(len(b_c_[0])):
            q_temp = self.target_net(b_g_, b_f_, b_c_[:, i, :].unsqueeze(1), b_hx_, b_r).detach()
            q_next[:, i] = q_temp.squeeze(dim=1)
        for j in range(len(b_c_)):
            if b_r[j] >= 1000:
                q_next[j, :] = torch.zeros((1, len(b_c_[0]))).to(device)
        q_target = b_r + GAMMA * q_next.max(1)[0].view(batch_size, 1)   # shape (batch, 1)
        # print("q_eval, q_target", q_eval, q_target, b_r)
        critic_loss = torch.mean(F.mse_loss(q_eval.float(), q_target.float()))

        # loss = self.loss_func(q_eval.float(), q_target.float())
        # self.optimizer.zero_grad()
        # loss.backward()
        # self.optimizer.step()

        log_probs = self.actor_net(b_g, b_f.detach(), b_c, b_hx.detach(), b_r)
        delta = q_target - q_eval
        # print("log_probs", log_probs)
        actor_loss = torch.mean( - log_probs * delta.detach()).float()

        self.actor_optimizer.zero_grad()
        actor_loss.backward()
        self.actor_optimizer.step()

        self.critic_optimizer.zero_grad()
        critic_loss.backward()
        self.critic_optimizer.step()

        return actor_loss, critic_loss