<reponame>zhaohengyin/irgail_example<gh_stars>1-10
import os
import numpy as np
import torch


class SerializedBuffer:

    def __init__(self, path, device):
        tmp = torch.load(path)
        self.buffer_size = self._n = tmp['state'].size(0)
        self.device = device

        self.states = tmp['state'].clone().to(self.device)
        self.next_actions = tmp['next_action'].clone().to(self.device)
        self.actions = tmp['action'].clone().to(self.device)
        self.rewards = tmp['reward'].clone().to(self.device)
        self.dones = tmp['done'].clone().to(self.device)
        self.next_states = tmp['next_state'].clone().to(self.device)
        self.last_states = tmp['last_state'].clone().to(self.device)
        self.last_actions = tmp['last_action'].clone().to(self.device)
        if 'code' in tmp:
            self.codes = tmp['code'].clone().to(self.device)
        else:
            self.codes = None

        self.prob = np.ones(self.states.size(0)) / self.states.size(0)

    def sample(self, batch_size):
        idxes = np.random.randint(low=0, high=self._n, size=batch_size)
        if self.codes is None:
            return (
                self.states[idxes],
                self.actions[idxes],
                self.rewards[idxes],
                self.dones[idxes],
                self.next_states[idxes],
                None
            )
        else:
            return (
                self.states[idxes],
                self.actions[idxes],
                self.rewards[idxes],
                self.dones[idxes],
                self.next_states[idxes],
                self.codes[idxes]
            )

    def sample_full(self, batch_size):
        idxes = np.random.randint(low=0, high=self._n, size=batch_size)
        return (
            self.states[idxes],
            self.actions[idxes],
            self.rewards[idxes],
            self.dones[idxes],
            self.last_actions[idxes],
            self.last_states[idxes],
            self.next_actions[idxes],
            self.next_states[idxes]
        )

    def sample_adaptive(self, batch_size):
        # We need to sample more data from the relevant
        idxes = np.random.choice(self.states.size(0), size=batch_size, p=self.prob)
        return (
            self.states[idxes],
            self.actions[idxes],
            self.rewards[idxes],
            self.dones[idxes],
            self.last_actions[idxes],
            self.last_states[idxes],
            self.next_actions[idxes],
            self.next_states[idxes]
        )

    def sample_mixture(self, batch_size):
        # We need to sample more data from the relevant
        idxes_0 = np.random.choice(self.states.size(0), size=batch_size // 2, p=self.prob)
        idxes_1 = np.random.randint(low=0, high=self._n, size=batch_size // 2)
        idxes = np.concatenate([idxes_0, idxes_1])
        return (
            self.states[idxes],
            self.actions[idxes],
            self.rewards[idxes],
            self.dones[idxes],
            self.last_actions[idxes],
            self.last_states[idxes],
            self.next_actions[idxes],
            self.next_states[idxes]
        )

    def update_prob(self, discriminator):
        probablities = discriminator.calculate_probabilities(self.states).cpu().detach().numpy()
        probablities = probablities.reshape(-1)
        probablities = probablities / np.sum(probablities)
        print(probablities.max(), probablities.min())
       # print(probablities.sum(), self.prob.sum())
        self.prob = 0.9 * probablities + 0.1 * self.prob
        self.prob = self.prob / np.sum(self.prob)
        print(self.prob.max(), self.prob.min())
        #print(self.prob.sum())

class SerializedTransformedBuffer:

    def __init__(self, path, device):
        tmp = torch.load(path)
        self.buffer_size = self._n = tmp['state'].size(0)
        self.device = device

        self.states = tmp['state'].clone().to(self.device)
        self.actions = tmp['action'].clone().to(self.device)
        self.rewards = tmp['reward'].clone().to(self.device)
        self.dones = tmp['done'].clone().to(self.device)
        self.next_states = tmp['next_state'].clone().to(self.device)
        # self.transform = torch.from_numpy(np.array([0.5, 0.2, -0.3, 0.5], np.float32)).to(self.device)

    def transform(self, states):
        a, b = states[:, 0:2], states[:, 2:]
        return torch.cat((a /2, b * 2), dim=1)

    def sample(self, batch_size):
        idxes = np.random.randint(low=0, high=self._n, size=batch_size)
        # print(self.states[idxes])

        return (
            self.transform(self.states[idxes]), #+ self.transform,
            self.actions[idxes],
            self.rewards[idxes],
            self.dones[idxes],
            self.transform(self.next_states[idxes]),# + self.transform
        )


class Buffer(SerializedBuffer):
    def __init__(self, buffer_size, state_shape, action_shape, device, code_shape=0):
        self._n = 0
        self._p = 0
        self.buffer_size = buffer_size
        self.device = device

        self.states = torch.empty(
            (buffer_size, *state_shape), dtype=torch.float, device=device)
        self.actions = torch.empty(
            (buffer_size, *action_shape), dtype=torch.float, device=device)

        self.last_states = torch.empty(
            (buffer_size, *state_shape), dtype=torch.float, device=device)
        self.last_actions = torch.empty(
            (buffer_size, *action_shape), dtype=torch.float, device=device)

        self.next_actions = torch.empty(
            (buffer_size, *action_shape), dtype=torch.float, device=device)

        self.rewards = torch.empty(
            (buffer_size, 1), dtype=torch.float, device=device)
        self.dones = torch.empty(
            (buffer_size, 1), dtype=torch.float, device=device)

        self.next_states = torch.empty(
            (buffer_size, *state_shape), dtype=torch.float, device=device)

        if code_shape != 0:
            self.codes =  torch.empty(
                (buffer_size, code_shape), dtype=torch.float, device=device)
        else:
            self.codes = None

    def append(self, state, action, reward, done, next_state, next_action=None, last_state=None, last_action=None, code=None):
        self.states[self._p].copy_(torch.from_numpy(state))
        self.actions[self._p].copy_(torch.from_numpy(action))
        self.rewards[self._p] = float(reward)
        self.dones[self._p] = float(done)
        self.next_states[self._p].copy_(torch.from_numpy(next_state))

        if next_action is not None:
            self.next_actions[self._p].copy_(torch.from_numpy(next_action))
        if last_action is not None:
            self.last_actions[self._p].copy_(torch.from_numpy(last_action))
        if last_state is not None:
            self.last_states[self._p].copy_(torch.from_numpy(last_state))
        if code is not None and self.codes is not None:
            self.codes[self._p].copy_(torch.from_numpy(code))

        self._p = (self._p + 1) % self.buffer_size
        self._n = min(self._n + 1, self.buffer_size)

    def save(self, path):
        if not os.path.exists(os.path.dirname(path)):
            os.makedirs(os.path.dirname(path))

        if self.codes is not None:
            torch.save({
                'state': self.states.clone().cpu(),
                'action': self.actions.clone().cpu(),
                'next_action': self.actions.clone().cpu(),
                'reward': self.rewards.clone().cpu(),
                'done': self.dones.clone().cpu(),
                'next_state': self.next_states.clone().cpu(),
                'last_state' :self.last_states.clone().cpu(),
                'last_action':self.last_actions.clone().cpu(),
                'code': self.codes.clone().cpu()
            }, path)
        else:
            torch.save({
                'state': self.states.clone().cpu(),
                'action': self.actions.clone().cpu(),
                'next_action': self.actions.clone().cpu(),
                'reward': self.rewards.clone().cpu(),
                'done': self.dones.clone().cpu(),
                'next_state': self.next_states.clone().cpu(),
                'last_state': self.last_states.clone().cpu(),
                'last_action': self.last_actions.clone().cpu()
            }, path)



class RolloutBuffer:

    def __init__(self, buffer_size, state_shape, action_shape, device, mix=1):
        self._n = 0
        self._p = 0
        self.mix = mix
        self.buffer_size = buffer_size
        self.total_size = mix * buffer_size

        self.states = torch.empty(
            (self.total_size, *state_shape), dtype=torch.float, device=device)
        self.actions = torch.empty(
            (self.total_size, *action_shape), dtype=torch.float, device=device)
        self.rewards = torch.empty(
            (self.total_size, 1), dtype=torch.float, device=device)
        self.dones = torch.empty(
            (self.total_size, 1), dtype=torch.float, device=device)
        self.log_pis = torch.empty(
            (self.total_size, 1), dtype=torch.float, device=device)
        self.next_states = torch.empty(
            (self.total_size, *state_shape), dtype=torch.float, device=device)

    def append(self, state, action, reward, done, log_pi, next_state):
        self.states[self._p].copy_(torch.from_numpy(state))
        self.actions[self._p].copy_(torch.from_numpy(action))
        self.rewards[self._p] = float(reward)
        self.dones[self._p] = float(done)
        self.log_pis[self._p] = float(log_pi)
        self.next_states[self._p].copy_(torch.from_numpy(next_state))

        self._p = (self._p + 1) % self.total_size
        self._n = min(self._n + 1, self.total_size)

    def get(self):
        assert self._p % self.buffer_size == 0
        start = (self._p - self.buffer_size) % self.total_size
        idxes = slice(start, start + self.buffer_size)
        return (
            self.states[idxes],
            self.actions[idxes],
            self.rewards[idxes],
            self.dones[idxes],
            self.log_pis[idxes],
            self.next_states[idxes]
        )

    def sample(self, batch_size):
        assert self._p % self.buffer_size == 0
        idxes = np.random.randint(low=0, high=self._n, size=batch_size)
        return (
            self.states[idxes],
            self.actions[idxes],
            self.rewards[idxes],
            self.dones[idxes],
            self.log_pis[idxes],
            self.next_states[idxes]
        )

import os
import numpy as np
import torch


class SerializedBuffer:

    def __init__(self, path, device):
        tmp = torch.load(path)
        self.buffer_size = self._n = tmp['state'].size(0)
        self.device = device

        self.states = tmp['state'].clone().to(self.device)
        self.next_actions = tmp['next_action'].clone().to(self.device)
        self.actions = tmp['action'].clone().to(self.device)
        self.rewards = tmp['reward'].clone().to(self.device)
        self.dones = tmp['done'].clone().to(self.device)
        self.next_states = tmp['next_state'].clone().to(self.device)
        self.last_states = tmp['last_state'].clone().to(self.device)
        self.last_actions = tmp['last_action'].clone().to(self.device)
        if 'code' in tmp:
            self.codes = tmp['code'].clone().to(self.device)
        else:
            self.codes = None

        self.prob = np.ones(self.states.size(0)) / self.states.size(0)

    def sample(self, batch_size):
        idxes = np.random.randint(low=0, high=self._n, size=batch_size)
        if self.codes is None:
            return (
                self.states[idxes],
                self.actions[idxes],
                self.rewards[idxes],
                self.dones[idxes],
                self.next_states[idxes],
                None
            )
        else:
            return (
                self.states[idxes],
                self.actions[idxes],
                self.rewards[idxes],
                self.dones[idxes],
                self.next_states[idxes],
                self.codes[idxes]
            )

    def sample_full(self, batch_size):
        idxes = np.random.randint(low=0, high=self._n, size=batch_size)
        return (
            self.states[idxes],
            self.actions[idxes],
            self.rewards[idxes],
            self.dones[idxes],
            self.last_actions[idxes],
            self.last_states[idxes],
            self.next_actions[idxes],
            self.next_states[idxes]
        )

    def sample_adaptive(self, batch_size):
        # We need to sample more data from the relevant
        idxes = np.random.choice(self.states.size(0), size=batch_size, p=self.prob)
        return (
            self.states[idxes],
            self.actions[idxes],
            self.rewards[idxes],
            self.dones[idxes],
            self.last_actions[idxes],
            self.last_states[idxes],
            self.next_actions[idxes],
            self.next_states[idxes]
        )

    def sample_mixture(self, batch_size):
        # We need to sample more data from the relevant
        idxes_0 = np.random.choice(self.states.size(0), size=batch_size // 2, p=self.prob)
        idxes_1 = np.random.randint(low=0, high=self._n, size=batch_size // 2)
        idxes = np.concatenate([idxes_0, idxes_1])
        return (
            self.states[idxes],
            self.actions[idxes],
            self.rewards[idxes],
            self.dones[idxes],
            self.last_actions[idxes],
            self.last_states[idxes],
            self.next_actions[idxes],
            self.next_states[idxes]
        )

    def update_prob(self, discriminator):
        probablities = discriminator.calculate_probabilities(self.states).cpu().detach().numpy()
        probablities = probablities.reshape(-1)
        probablities = probablities / np.sum(probablities)
        print(probablities.max(), probablities.min())
       # print(probablities.sum(), self.prob.sum())
        self.prob = 0.9 * probablities + 0.1 * self.prob
        self.prob = self.prob / np.sum(self.prob)
        print(self.prob.max(), self.prob.min())
        #print(self.prob.sum())

class SerializedTransformedBuffer:

    def __init__(self, path, device):
        tmp = torch.load(path)
        self.buffer_size = self._n = tmp['state'].size(0)
        self.device = device

        self.states = tmp['state'].clone().to(self.device)
        self.actions = tmp['action'].clone().to(self.device)
        self.rewards = tmp['reward'].clone().to(self.device)
        self.dones = tmp['done'].clone().to(self.device)
        self.next_states = tmp['next_state'].clone().to(self.device)
        # self.transform = torch.from_numpy(np.array([0.5, 0.2, -0.3, 0.5], np.float32)).to(self.device)

    def transform(self, states):
        a, b = states[:, 0:2], states[:, 2:]
        return torch.cat((a /2, b * 2), dim=1)

    def sample(self, batch_size):
        idxes = np.random.randint(low=0, high=self._n, size=batch_size)
        # print(self.states[idxes])

        return (
            self.transform(self.states[idxes]), #+ self.transform,
            self.actions[idxes],
            self.rewards[idxes],
            self.dones[idxes],
            self.transform(self.next_states[idxes]),# + self.transform
        )


class Buffer(SerializedBuffer):
    def __init__(self, buffer_size, state_shape, action_shape, device, code_shape=0):
        self._n = 0
        self._p = 0
        self.buffer_size = buffer_size
        self.device = device

        self.states = torch.empty(
            (buffer_size, *state_shape), dtype=torch.float, device=device)
        self.actions = torch.empty(
            (buffer_size, *action_shape), dtype=torch.float, device=device)

        self.last_states = torch.empty(
            (buffer_size, *state_shape), dtype=torch.float, device=device)
        self.last_actions = torch.empty(
            (buffer_size, *action_shape), dtype=torch.float, device=device)

        self.next_actions = torch.empty(
            (buffer_size, *action_shape), dtype=torch.float, device=device)

        self.rewards = torch.empty(
            (buffer_size, 1), dtype=torch.float, device=device)
        self.dones = torch.empty(
            (buffer_size, 1), dtype=torch.float, device=device)

        self.next_states = torch.empty(
            (buffer_size, *state_shape), dtype=torch.float, device=device)

        if code_shape != 0:
            self.codes =  torch.empty(
                (buffer_size, code_shape), dtype=torch.float, device=device)
        else:
            self.codes = None

    def append(self, state, action, reward, done, next_state, next_action=None, last_state=None, last_action=None, code=None):
        self.states[self._p].copy_(torch.from_numpy(state))
        self.actions[self._p].copy_(torch.from_numpy(action))
        self.rewards[self._p] = float(reward)
        self.dones[self._p] = float(done)
        self.next_states[self._p].copy_(torch.from_numpy(next_state))

        if next_action is not None:
            self.next_actions[self._p].copy_(torch.from_numpy(next_action))
        if last_action is not None:
            self.last_actions[self._p].copy_(torch.from_numpy(last_action))
        if last_state is not None:
            self.last_states[self._p].copy_(torch.from_numpy(last_state))
        if code is not None and self.codes is not None:
            self.codes[self._p].copy_(torch.from_numpy(code))

        self._p = (self._p + 1) % self.buffer_size
        self._n = min(self._n + 1, self.buffer_size)

    def save(self, path):
        if not os.path.exists(os.path.dirname(path)):
            os.makedirs(os.path.dirname(path))

        if self.codes is not None:
            torch.save({
                'state': self.states.clone().cpu(),
                'action': self.actions.clone().cpu(),
                'next_action': self.actions.clone().cpu(),
                'reward': self.rewards.clone().cpu(),
                'done': self.dones.clone().cpu(),
                'next_state': self.next_states.clone().cpu(),
                'last_state' :self.last_states.clone().cpu(),
                'last_action':self.last_actions.clone().cpu(),
                'code': self.codes.clone().cpu()
            }, path)
        else:
            torch.save({
                'state': self.states.clone().cpu(),
                'action': self.actions.clone().cpu(),
                'next_action': self.actions.clone().cpu(),
                'reward': self.rewards.clone().cpu(),
                'done': self.dones.clone().cpu(),
                'next_state': self.next_states.clone().cpu(),
                'last_state': self.last_states.clone().cpu(),
                'last_action': self.last_actions.clone().cpu()
            }, path)



class RolloutBuffer:

    def __init__(self, buffer_size, state_shape, action_shape, device, mix=1):
        self._n = 0
        self._p = 0
        self.mix = mix
        self.buffer_size = buffer_size
        self.total_size = mix * buffer_size

        self.states = torch.empty(
            (self.total_size, *state_shape), dtype=torch.float, device=device)
        self.actions = torch.empty(
            (self.total_size, *action_shape), dtype=torch.float, device=device)
        self.rewards = torch.empty(
            (self.total_size, 1), dtype=torch.float, device=device)
        self.dones = torch.empty(
            (self.total_size, 1), dtype=torch.float, device=device)
        self.log_pis = torch.empty(
            (self.total_size, 1), dtype=torch.float, device=device)
        self.next_states = torch.empty(
            (self.total_size, *state_shape), dtype=torch.float, device=device)

    def append(self, state, action, reward, done, log_pi, next_state):
        self.states[self._p].copy_(torch.from_numpy(state))
        self.actions[self._p].copy_(torch.from_numpy(action))
        self.rewards[self._p] = float(reward)
        self.dones[self._p] = float(done)
        self.log_pis[self._p] = float(log_pi)
        self.next_states[self._p].copy_(torch.from_numpy(next_state))

        self._p = (self._p + 1) % self.total_size
        self._n = min(self._n + 1, self.total_size)

    def get(self):
        assert self._p % self.buffer_size == 0
        start = (self._p - self.buffer_size) % self.total_size
        idxes = slice(start, start + self.buffer_size)
        return (
            self.states[idxes],
            self.actions[idxes],
            self.rewards[idxes],
            self.dones[idxes],
            self.log_pis[idxes],
            self.next_states[idxes]
        )

    def sample(self, batch_size):
        assert self._p % self.buffer_size == 0
        idxes = np.random.randint(low=0, high=self._n, size=batch_size)
        return (
            self.states[idxes],
            self.actions[idxes],
            self.rewards[idxes],
            self.dones[idxes],
            self.log_pis[idxes],
            self.next_states[idxes]
        )
# We need to sample more data from the relevant # We need to sample more data from the relevant # print(probablities.sum(), self.prob.sum()) #print(self.prob.sum()) # self.transform = torch.from_numpy(np.array([0.5, 0.2, -0.3, 0.5], np.float32)).to(self.device) # print(self.states[idxes]) #+ self.transform, # + self.transform