# MIT License
#
# Copyright (C) 2021. Huawei Technologies Co., Ltd. All rights reserved.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
# Do not make any change to this file when merging. Just use my version.
from collections import deque, namedtuple
import numpy as np
import random, copy
import torch
from marl_scalability.utils.common import normalize_im
from collections.abc import Iterable

from torch.utils.data import Dataset, Sampler, DataLoader

Transition = namedtuple(
    "Transition",
    field_names=["state", "action", "reward", "next_state", "done", "others"],
    # others may contain importance sampling ratio, GVF rewards,... etc
    defaults=(None,) * 6,
)


class RandomRLSampler(Sampler):
    def __init__(self, data_source, batch_size):
        self.datasource = data_source
        self.batch_size = batch_size

    def __len__(self):
        return len(self.datasource)

    def __iter__(self):
        n = len(self.datasource)
        return iter(torch.randperm(n).tolist()[0 : self.batch_size])


class ReplayBufferDataset(Dataset):
    cpu = torch.device("cpu")

    def __init__(self, buffer_size, device):
        self.buffer_size = buffer_size
        self.memory = deque(maxlen=self.buffer_size)
        self.device = device

    def add(
        self,
        state,
        action,
        reward,
        next_state,
        done,
        prev_action,
        others=None,
    ):
        if others is None:
            others = {}
        # dereference the states
        state = copy.deepcopy(state)
        next_state = copy.deepcopy(next_state)
        state["low_dim_states"] = np.float32(
            np.append(state["low_dim_states"], prev_action)
        )
        state["low_dim_states"] = torch.from_numpy(state["low_dim_states"]).to(
            self.device
        )
        state["social_vehicles"] = torch.from_numpy(state["social_vehicles"]).to(
            self.device
        )

        next_state["low_dim_states"] = np.float32(
            np.append(next_state["low_dim_states"], action)
        )
        next_state["social_vehicles"] = torch.from_numpy(
            next_state["social_vehicles"]
        ).to(self.device)
        next_state["low_dim_states"] = torch.from_numpy(
            next_state["low_dim_states"]
        ).to(self.device)

        action = np.asarray([action]) if not isinstance(action, Iterable) else action
        action = torch.from_numpy(action).float()
        reward = torch.from_numpy(np.asarray([reward])).float()
        done = torch.from_numpy(np.asarray([done])).float()
        new_experience = Transition(state, action, reward, next_state, done, others)
        self.memory.append(new_experience)

    def __len__(self):
        return len(self.memory)

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()
        state, action, reward, next_state, done, others = tuple(self.memory[idx])
        return state, action, reward, next_state, done, others


class ReplayBuffer:
    def __init__(
        self,
        buffer_size,
        batch_size,
        device_name,
        pin_memory=False,
        num_workers=0,
    ):
        self.replay_buffer_dataset = ReplayBufferDataset(buffer_size, device=None)
        self.sampler = RandomRLSampler(self.replay_buffer_dataset, batch_size)
        self.data_loader = DataLoader(
            self.replay_buffer_dataset,
            sampler=self.sampler,
            pin_memory=pin_memory,
            num_workers=num_workers,
        )
        self.storage_device = torch.device(device_name)

    def add(self, *args, **kwargs):
        self.replay_buffer_dataset.add(*args, **kwargs)

    def __len__(self):
        return len(self.replay_buffer_dataset)

    def __getitem__(self, idx):
        return self.replay_buffer_dataset[idx]

    def make_state_from_dict(self, states, device):
        # image_keys = states[0]["images"].keys()
        # images = {}
        # for k in image_keys:
        #     _images = torch.cat([e[k] for e in states], dim=0).float().to(device)
        #     _images = normalize_im(_images)
        #     images[k] = _images
        low_dim_states = (
            torch.cat([e["low_dim_states"] for e in states], dim=0).float().to(device)
        )
        if "social_vehicles" in states[0]:
            social_vehicles = [
                e["social_vehicles"][0].float().to(device) for e in states
            ]
        else:
            social_vehicles = False
        out = {
            # "images": images,
            "low_dim_states": low_dim_states,
            "social_vehicles": social_vehicles,
        }
        return out

    def sample(self, device=None):
        device = device if device else self.storage_device
        batch = list(iter(self.data_loader))
        states, actions, rewards, next_states, dones, others = zip(*batch)
        states = self.make_state_from_dict(states, device)
        next_states = self.make_state_from_dict(next_states, device)
        actions = torch.cat(actions, dim=0).float().to(device)
        rewards = torch.cat(rewards, dim=0).float().to(device)
        dones = torch.cat(dones, dim=0).float().to(device)
        return states, actions, rewards, next_states, dones, others

ebd53bfa911aa0b65c839ebcc276e1c351f420b0
<mask token>


class ReplayBufferDataset(Dataset):
    <mask token>

    def __init__(self, buffer_size, device):
        self.buffer_size = buffer_size
        self.memory = deque(maxlen=self.buffer_size)
        self.device = device

    def add(self, state, action, reward, next_state, done, prev_action,
        others=None):
        if others is None:
            others = {}
        state = copy.deepcopy(state)
        next_state = copy.deepcopy(next_state)
        state['low_dim_states'] = np.float32(np.append(state[
            'low_dim_states'], prev_action))
        state['low_dim_states'] = torch.from_numpy(state['low_dim_states']).to(
            self.device)
        state['social_vehicles'] = torch.from_numpy(state['social_vehicles']
            ).to(self.device)
        next_state['low_dim_states'] = np.float32(np.append(next_state[
            'low_dim_states'], action))
        next_state['social_vehicles'] = torch.from_numpy(next_state[
            'social_vehicles']).to(self.device)
        next_state['low_dim_states'] = torch.from_numpy(next_state[
            'low_dim_states']).to(self.device)
        action = np.asarray([action]) if not isinstance(action, Iterable
            ) else action
        action = torch.from_numpy(action).float()
        reward = torch.from_numpy(np.asarray([reward])).float()
        done = torch.from_numpy(np.asarray([done])).float()
        new_experience = Transition(state, action, reward, next_state, done,
            others)
        self.memory.append(new_experience)

    def __len__(self):
        return len(self.memory)

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()
        state, action, reward, next_state, done, others = tuple(self.memory
            [idx])
        return state, action, reward, next_state, done, others


class ReplayBuffer:

    def __init__(self, buffer_size, batch_size, device_name, pin_memory=
        False, num_workers=0):
        self.replay_buffer_dataset = ReplayBufferDataset(buffer_size,
            device=None)
        self.sampler = RandomRLSampler(self.replay_buffer_dataset, batch_size)
        self.data_loader = DataLoader(self.replay_buffer_dataset, sampler=
            self.sampler, pin_memory=pin_memory, num_workers=num_workers)
        self.storage_device = torch.device(device_name)

    def add(self, *args, **kwargs):
        self.replay_buffer_dataset.add(*args, **kwargs)

    def __len__(self):
        return len(self.replay_buffer_dataset)

    def __getitem__(self, idx):
        return self.replay_buffer_dataset[idx]

    def make_state_from_dict(self, states, device):
        low_dim_states = torch.cat([e['low_dim_states'] for e in states], dim=0
            ).float().to(device)
        if 'social_vehicles' in states[0]:
            social_vehicles = [e['social_vehicles'][0].float().to(device) for
                e in states]
        else:
            social_vehicles = False
        out = {'low_dim_states': low_dim_states, 'social_vehicles':
            social_vehicles}
        return out

    def sample(self, device=None):
        device = device if device else self.storage_device
        batch = list(iter(self.data_loader))
        states, actions, rewards, next_states, dones, others = zip(*batch)
        states = self.make_state_from_dict(states, device)
        next_states = self.make_state_from_dict(next_states, device)
        actions = torch.cat(actions, dim=0).float().to(device)
        rewards = torch.cat(rewards, dim=0).float().to(device)
        dones = torch.cat(dones, dim=0).float().to(device)
        return states, actions, rewards, next_states, dones, others

<mask token>


class ReplayBufferDataset(Dataset):
    cpu = torch.device('cpu')

    def __init__(self, buffer_size, device):
        self.buffer_size = buffer_size
        self.memory = deque(maxlen=self.buffer_size)
        self.device = device

    def add(self, state, action, reward, next_state, done, prev_action,
        others=None):
        if others is None:
            others = {}
        state = copy.deepcopy(state)
        next_state = copy.deepcopy(next_state)
        state['low_dim_states'] = np.float32(np.append(state[
            'low_dim_states'], prev_action))
        state['low_dim_states'] = torch.from_numpy(state['low_dim_states']).to(
            self.device)
        state['social_vehicles'] = torch.from_numpy(state['social_vehicles']
            ).to(self.device)
        next_state['low_dim_states'] = np.float32(np.append(next_state[
            'low_dim_states'], action))
        next_state['social_vehicles'] = torch.from_numpy(next_state[
            'social_vehicles']).to(self.device)
        next_state['low_dim_states'] = torch.from_numpy(next_state[
            'low_dim_states']).to(self.device)
        action = np.asarray([action]) if not isinstance(action, Iterable
            ) else action
        action = torch.from_numpy(action).float()
        reward = torch.from_numpy(np.asarray([reward])).float()
        done = torch.from_numpy(np.asarray([done])).float()
        new_experience = Transition(state, action, reward, next_state, done,
            others)
        self.memory.append(new_experience)

    def __len__(self):
        return len(self.memory)

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()
        state, action, reward, next_state, done, others = tuple(self.memory
            [idx])
        return state, action, reward, next_state, done, others


class ReplayBuffer:

    def __init__(self, buffer_size, batch_size, device_name, pin_memory=
        False, num_workers=0):
        self.replay_buffer_dataset = ReplayBufferDataset(buffer_size,
            device=None)
        self.sampler = RandomRLSampler(self.replay_buffer_dataset, batch_size)
        self.data_loader = DataLoader(self.replay_buffer_dataset, sampler=
            self.sampler, pin_memory=pin_memory, num_workers=num_workers)
        self.storage_device = torch.device(device_name)

    def add(self, *args, **kwargs):
        self.replay_buffer_dataset.add(*args, **kwargs)

    def __len__(self):
        return len(self.replay_buffer_dataset)

    def __getitem__(self, idx):
        return self.replay_buffer_dataset[idx]

    def make_state_from_dict(self, states, device):
        low_dim_states = torch.cat([e['low_dim_states'] for e in states], dim=0
            ).float().to(device)
        if 'social_vehicles' in states[0]:
            social_vehicles = [e['social_vehicles'][0].float().to(device) for
                e in states]
        else:
            social_vehicles = False
        out = {'low_dim_states': low_dim_states, 'social_vehicles':
            social_vehicles}
        return out

    def sample(self, device=None):
        device = device if device else self.storage_device
        batch = list(iter(self.data_loader))
        states, actions, rewards, next_states, dones, others = zip(*batch)
        states = self.make_state_from_dict(states, device)
        next_states = self.make_state_from_dict(next_states, device)
        actions = torch.cat(actions, dim=0).float().to(device)
        rewards = torch.cat(rewards, dim=0).float().to(device)
        dones = torch.cat(dones, dim=0).float().to(device)
        return states, actions, rewards, next_states, dones, others

<mask token>
Transition = namedtuple('Transition', field_names=['state', 'action',
    'reward', 'next_state', 'done', 'others'], defaults=(None,) * 6)


class RandomRLSampler(Sampler):

    def __init__(self, data_source, batch_size):
        self.datasource = data_source
        self.batch_size = batch_size

    def __len__(self):
        return len(self.datasource)

    def __iter__(self):
        n = len(self.datasource)
        return iter(torch.randperm(n).tolist()[0:self.batch_size])


class ReplayBufferDataset(Dataset):
    cpu = torch.device('cpu')

    def __init__(self, buffer_size, device):
        self.buffer_size = buffer_size
        self.memory = deque(maxlen=self.buffer_size)
        self.device = device

    def add(self, state, action, reward, next_state, done, prev_action,
        others=None):
        if others is None:
            others = {}
        state = copy.deepcopy(state)
        next_state = copy.deepcopy(next_state)
        state['low_dim_states'] = np.float32(np.append(state[
            'low_dim_states'], prev_action))
        state['low_dim_states'] = torch.from_numpy(state['low_dim_states']).to(
            self.device)
        state['social_vehicles'] = torch.from_numpy(state['social_vehicles']
            ).to(self.device)
        next_state['low_dim_states'] = np.float32(np.append(next_state[
            'low_dim_states'], action))
        next_state['social_vehicles'] = torch.from_numpy(next_state[
            'social_vehicles']).to(self.device)
        next_state['low_dim_states'] = torch.from_numpy(next_state[
            'low_dim_states']).to(self.device)
        action = np.asarray([action]) if not isinstance(action, Iterable
            ) else action
        action = torch.from_numpy(action).float()
        reward = torch.from_numpy(np.asarray([reward])).float()
        done = torch.from_numpy(np.asarray([done])).float()
        new_experience = Transition(state, action, reward, next_state, done,
            others)
        self.memory.append(new_experience)

    def __len__(self):
        return len(self.memory)

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()
        state, action, reward, next_state, done, others = tuple(self.memory
            [idx])
        return state, action, reward, next_state, done, others


class ReplayBuffer:

    def __init__(self, buffer_size, batch_size, device_name, pin_memory=
        False, num_workers=0):
        self.replay_buffer_dataset = ReplayBufferDataset(buffer_size,
            device=None)
        self.sampler = RandomRLSampler(self.replay_buffer_dataset, batch_size)
        self.data_loader = DataLoader(self.replay_buffer_dataset, sampler=
            self.sampler, pin_memory=pin_memory, num_workers=num_workers)
        self.storage_device = torch.device(device_name)

    def add(self, *args, **kwargs):
        self.replay_buffer_dataset.add(*args, **kwargs)

    def __len__(self):
        return len(self.replay_buffer_dataset)

    def __getitem__(self, idx):
        return self.replay_buffer_dataset[idx]

    def make_state_from_dict(self, states, device):
        low_dim_states = torch.cat([e['low_dim_states'] for e in states], dim=0
            ).float().to(device)
        if 'social_vehicles' in states[0]:
            social_vehicles = [e['social_vehicles'][0].float().to(device) for
                e in states]
        else:
            social_vehicles = False
        out = {'low_dim_states': low_dim_states, 'social_vehicles':
            social_vehicles}
        return out

    def sample(self, device=None):
        device = device if device else self.storage_device
        batch = list(iter(self.data_loader))
        states, actions, rewards, next_states, dones, others = zip(*batch)
        states = self.make_state_from_dict(states, device)
        next_states = self.make_state_from_dict(next_states, device)
        actions = torch.cat(actions, dim=0).float().to(device)
        rewards = torch.cat(rewards, dim=0).float().to(device)
        dones = torch.cat(dones, dim=0).float().to(device)
        return states, actions, rewards, next_states, dones, others

from collections import deque, namedtuple
import numpy as np
import random, copy
import torch
from marl_scalability.utils.common import normalize_im
from collections.abc import Iterable
from torch.utils.data import Dataset, Sampler, DataLoader
Transition = namedtuple('Transition', field_names=['state', 'action',
    'reward', 'next_state', 'done', 'others'], defaults=(None,) * 6)


class RandomRLSampler(Sampler):

    def __init__(self, data_source, batch_size):
        self.datasource = data_source
        self.batch_size = batch_size

    def __len__(self):
        return len(self.datasource)

    def __iter__(self):
        n = len(self.datasource)
        return iter(torch.randperm(n).tolist()[0:self.batch_size])


class ReplayBufferDataset(Dataset):
    cpu = torch.device('cpu')

    def __init__(self, buffer_size, device):
        self.buffer_size = buffer_size
        self.memory = deque(maxlen=self.buffer_size)
        self.device = device

    def add(self, state, action, reward, next_state, done, prev_action,
        others=None):
        if others is None:
            others = {}
        state = copy.deepcopy(state)
        next_state = copy.deepcopy(next_state)
        state['low_dim_states'] = np.float32(np.append(state[
            'low_dim_states'], prev_action))
        state['low_dim_states'] = torch.from_numpy(state['low_dim_states']).to(
            self.device)
        state['social_vehicles'] = torch.from_numpy(state['social_vehicles']
            ).to(self.device)
        next_state['low_dim_states'] = np.float32(np.append(next_state[
            'low_dim_states'], action))
        next_state['social_vehicles'] = torch.from_numpy(next_state[
            'social_vehicles']).to(self.device)
        next_state['low_dim_states'] = torch.from_numpy(next_state[
            'low_dim_states']).to(self.device)
        action = np.asarray([action]) if not isinstance(action, Iterable
            ) else action
        action = torch.from_numpy(action).float()
        reward = torch.from_numpy(np.asarray([reward])).float()
        done = torch.from_numpy(np.asarray([done])).float()
        new_experience = Transition(state, action, reward, next_state, done,
            others)
        self.memory.append(new_experience)

    def __len__(self):
        return len(self.memory)

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()
        state, action, reward, next_state, done, others = tuple(self.memory
            [idx])
        return state, action, reward, next_state, done, others


class ReplayBuffer:

    def __init__(self, buffer_size, batch_size, device_name, pin_memory=
        False, num_workers=0):
        self.replay_buffer_dataset = ReplayBufferDataset(buffer_size,
            device=None)
        self.sampler = RandomRLSampler(self.replay_buffer_dataset, batch_size)
        self.data_loader = DataLoader(self.replay_buffer_dataset, sampler=
            self.sampler, pin_memory=pin_memory, num_workers=num_workers)
        self.storage_device = torch.device(device_name)

    def add(self, *args, **kwargs):
        self.replay_buffer_dataset.add(*args, **kwargs)

    def __len__(self):
        return len(self.replay_buffer_dataset)

    def __getitem__(self, idx):
        return self.replay_buffer_dataset[idx]

    def make_state_from_dict(self, states, device):
        low_dim_states = torch.cat([e['low_dim_states'] for e in states], dim=0
            ).float().to(device)
        if 'social_vehicles' in states[0]:
            social_vehicles = [e['social_vehicles'][0].float().to(device) for
                e in states]
        else:
            social_vehicles = False
        out = {'low_dim_states': low_dim_states, 'social_vehicles':
            social_vehicles}
        return out

    def sample(self, device=None):
        device = device if device else self.storage_device
        batch = list(iter(self.data_loader))
        states, actions, rewards, next_states, dones, others = zip(*batch)
        states = self.make_state_from_dict(states, device)
        next_states = self.make_state_from_dict(next_states, device)
        actions = torch.cat(actions, dim=0).float().to(device)
        rewards = torch.cat(rewards, dim=0).float().to(device)
        dones = torch.cat(dones, dim=0).float().to(device)
        return states, actions, rewards, next_states, dones, others

# MIT License
#
# Copyright (C) 2021. Huawei Technologies Co., Ltd. All rights reserved.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
# Do not make any change to this file when merging. Just use my version.
from collections import deque, namedtuple
import numpy as np
import random, copy
import torch
from marl_scalability.utils.common import normalize_im
from collections.abc import Iterable

from torch.utils.data import Dataset, Sampler, DataLoader

Transition = namedtuple(
    "Transition",
    field_names=["state", "action", "reward", "next_state", "done", "others"],
    # others may contain importance sampling ratio, GVF rewards,... etc
    defaults=(None,) * 6,
)


class RandomRLSampler(Sampler):
    def __init__(self, data_source, batch_size):
        self.datasource = data_source
        self.batch_size = batch_size

    def __len__(self):
        return len(self.datasource)

    def __iter__(self):
        n = len(self.datasource)
        return iter(torch.randperm(n).tolist()[0 : self.batch_size])


class ReplayBufferDataset(Dataset):
    cpu = torch.device("cpu")

    def __init__(self, buffer_size, device):
        self.buffer_size = buffer_size
        self.memory = deque(maxlen=self.buffer_size)
        self.device = device

    def add(
        self,
        state,
        action,
        reward,
        next_state,
        done,
        prev_action,
        others=None,
    ):
        if others is None:
            others = {}
        # dereference the states
        state = copy.deepcopy(state)
        next_state = copy.deepcopy(next_state)
        state["low_dim_states"] = np.float32(
            np.append(state["low_dim_states"], prev_action)
        )
        state["low_dim_states"] = torch.from_numpy(state["low_dim_states"]).to(
            self.device
        )
        state["social_vehicles"] = torch.from_numpy(state["social_vehicles"]).to(
            self.device
        )

        next_state["low_dim_states"] = np.float32(
            np.append(next_state["low_dim_states"], action)
        )
        next_state["social_vehicles"] = torch.from_numpy(
            next_state["social_vehicles"]
        ).to(self.device)
        next_state["low_dim_states"] = torch.from_numpy(
            next_state["low_dim_states"]
        ).to(self.device)

        action = np.asarray([action]) if not isinstance(action, Iterable) else action
        action = torch.from_numpy(action).float()
        reward = torch.from_numpy(np.asarray([reward])).float()
        done = torch.from_numpy(np.asarray([done])).float()
        new_experience = Transition(state, action, reward, next_state, done, others)
        self.memory.append(new_experience)

    def __len__(self):
        return len(self.memory)

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()
        state, action, reward, next_state, done, others = tuple(self.memory[idx])
        return state, action, reward, next_state, done, others


class ReplayBuffer:
    def __init__(
        self,
        buffer_size,
        batch_size,
        device_name,
        pin_memory=False,
        num_workers=0,
    ):
        self.replay_buffer_dataset = ReplayBufferDataset(buffer_size, device=None)
        self.sampler = RandomRLSampler(self.replay_buffer_dataset, batch_size)
        self.data_loader = DataLoader(
            self.replay_buffer_dataset,
            sampler=self.sampler,
            pin_memory=pin_memory,
            num_workers=num_workers,
        )
        self.storage_device = torch.device(device_name)

    def add(self, *args, **kwargs):
        self.replay_buffer_dataset.add(*args, **kwargs)

    def __len__(self):
        return len(self.replay_buffer_dataset)

    def __getitem__(self, idx):
        return self.replay_buffer_dataset[idx]

    def make_state_from_dict(self, states, device):
        # image_keys = states[0]["images"].keys()
        # images = {}
        # for k in image_keys:
        #     _images = torch.cat([e[k] for e in states], dim=0).float().to(device)
        #     _images = normalize_im(_images)
        #     images[k] = _images
        low_dim_states = (
            torch.cat([e["low_dim_states"] for e in states], dim=0).float().to(device)
        )
        if "social_vehicles" in states[0]:
            social_vehicles = [
                e["social_vehicles"][0].float().to(device) for e in states
            ]
        else:
            social_vehicles = False
        out = {
            # "images": images,
            "low_dim_states": low_dim_states,
            "social_vehicles": social_vehicles,
        }
        return out

    def sample(self, device=None):
        device = device if device else self.storage_device
        batch = list(iter(self.data_loader))
        states, actions, rewards, next_states, dones, others = zip(*batch)
        states = self.make_state_from_dict(states, device)
        next_states = self.make_state_from_dict(next_states, device)
        actions = torch.cat(actions, dim=0).float().to(device)
        rewards = torch.cat(rewards, dim=0).float().to(device)
        dones = torch.cat(dones, dim=0).float().to(device)
        return states, actions, rewards, next_states, dones, others
