data: namedtuple,

gamma: float,

criterion: torch.nn.modules = nn.MSELoss(reduction='none') # noqa

) -> torch.Tensor:

q, next_q, act, next_act, reward, done, weight = data

assert len(act.shape) == 1, act.shape

assert len(reward.shape) == 1, reward.shape

batch_range = torch.arange(act.shape[0])

if weight is None:

weight = torch.ones_like(reward)

q_s_a = q[batch_range, act]

target_q_s_a = next_q[batch_range, next_act]

target_q_s_a = gamma * (1 - done) * target_q_s_a + reward

reward, next_value, done = data

assert reward.shape[0] == nstep

device = reward.device

reward_factor = torch.ones(nstep).to(device)

for i in range(1, nstep):

reward_factor[i] = gamma * reward_factor[i - 1]

reward = torch.matmul(reward_factor, reward)

if value_gamma is None:

return_ = reward + (gamma ** nstep) * next_value * (1 - done)

else:

return_ = reward + value_gamma * next_value * (1 - done)

data: namedtuple,

gamma: float,

v_min: float,

v_max: float,

n_atom: int,

) -> torch.Tensor:

dist, next_dist, act, next_act, reward, done, weight = data

device = reward.device

assert len(act.shape) == 1, act.shape

assert len(reward.shape) == 1, reward.shape

reward = reward.unsqueeze(-1)

done = done.unsqueeze(-1)

support = torch.linspace(v_min, v_max, n_atom).to(device)

delta_z = (v_max - v_min) / (n_atom - 1)

batch_size = act.shape[0]

batch_range = torch.arange(batch_size)

if weight is None:

weight = torch.ones_like(reward)

next_dist = next_dist[batch_range, next_act].detach()

target_z = reward + (1 - done) * gamma * support

target_z = target_z.clamp(min=v_min, max=v_max)

b = (target_z - v_min) / delta_z

l = b.floor().long()

u = b.ceil().long()

# Fix disappearing probability mass when l = b = u (b is int)

l[(u > 0) * (l == u)] -= 1

u[(l < (n_atom - 1)) * (l == u)] += 1

proj_dist = torch.zeros_like(next_dist)

offset = torch.linspace(0, (batch_size - 1) * n_atom, batch_size).unsqueeze(1).expand(batch_size,

n_atom).long().to(device)

proj_dist.view(-1).index_add_(0, (l + offset).view(-1), (next_dist * (u.float() - b)).view(-1))

proj_dist.view(-1).index_add_(0, (u + offset).view(-1), (next_dist * (b - l.float())).view(-1))

log_p = torch.log(dist[batch_range, act])

loss = -(log_p * proj_dist * weight).sum(-1).mean()

r"""

if len(args) <= 0:

tmp = [kwargs['data'].reward.shape[0]]

tmp.extend(list(kwargs['data'].dist.shape))

else:

tmp = [args[0].reward.shape[0]]

tmp.extend(list(args[0].dist.shape))

data: namedtuple,

gamma: float,

v_min: float,

v_max: float,

n_atom: int,

nstep: int = 1,

value_gamma: Optional[torch.Tensor] = None,

) -> torch.Tensor:

r"""

dist, next_n_dist, act, next_n_act, reward, done, weight = data

device = reward.device

assert len(act.shape) == 1, act.shape

reward_factor = torch.ones(nstep).to(device)

for i in range(1, nstep):

reward_factor[i] = gamma * reward_factor[i - 1]

reward = torch.matmul(reward_factor, reward)

reward = reward.unsqueeze(-1)

done = done.unsqueeze(-1)

support = torch.linspace(v_min, v_max, n_atom).to(device)

delta_z = (v_max - v_min) / (n_atom - 1)

batch_size = act.shape[0]

batch_range = torch.arange(batch_size)

if weight is None:

weight = torch.ones_like(reward)

next_n_dist = next_n_dist[batch_range, next_n_act].detach()

if value_gamma is None:

target_z = reward + (1 - done) * (gamma ** nstep) * support

else:

value_gamma = value_gamma.unsqueeze(-1)

target_z = reward + (1 - done) * value_gamma * support

target_z = target_z.clamp(min=v_min, max=v_max)

b = (target_z - v_min) / delta_z

l = b.floor().long()

u = b.ceil().long()

# Fix disappearing probability mass when l = b = u (b is int)

l[(u > 0) * (l == u)] -= 1

u[(l < (n_atom - 1)) * (l == u)] += 1

proj_dist = torch.zeros_like(next_n_dist)

offset = torch.linspace(0, (batch_size - 1) * n_atom, batch_size).unsqueeze(1).expand(batch_size,

n_atom).long().to(device)

proj_dist.view(-1).index_add_(0, (l + offset).view(-1), (next_n_dist * (u.float() - b)).view(-1))

proj_dist.view(-1).index_add_(0, (u + offset).view(-1), (next_n_dist * (b - l.float())).view(-1))

assert (dist[batch_range, act] > 0.0).all(), ("dist act", dist[batch_range, act], "dist:", dist)

log_p = torch.log(dist[batch_range, act])

if len(weight.shape) == 1:

weight = weight.unsqueeze(-1)

td_error_per_sample = -(log_p * proj_dist).sum(-1)

loss = -(log_p * proj_dist * weight).sum(-1).mean()

data: namedtuple,

gamma: float,

criterion: torch.nn.modules = nn.MSELoss(reduction='none') # noqa

) -> torch.Tensor:

v, next_v, reward, done, weight = data

if weight is None:

weight = torch.ones_like(reward)

if done is not None:

target_v = gamma * (1 - done) * next_v + reward

else:

target_v = gamma * next_v + reward

td_error_per_sample = criterion(v, target_v.detach())

data: namedtuple,

gamma: float,

nstep: int = 1,

criterion: torch.nn.modules = nn.MSELoss(reduction='none') # noqa

) -> torch.Tensor:

r"""

v, next_n_v, reward, done, weight, value_gamma = data

if weight is None:

weight = torch.ones_like(v)

target_v = nstep_return(nstep_return_data(reward, next_n_v, done), gamma, nstep, value_gamma)

td_error_per_sample = criterion(v, target_v.detach())

'dqfd_nstep_td_data', [

'q', 'next_n_q', 'action', 'next_n_action', 'reward', 'done', 'done_1', 'weight', 'new_n_q_one_step',

'next_n_action_one_step', 'is_expert'

]

r"""

if len(args) <= 0:

tmp = [kwargs['data'].reward.shape[0]]

tmp.extend(list(kwargs['data'].q.shape))

else:

tmp = [args[0].reward.shape[0]]

tmp.extend(list(args[0].q.shape))

data: namedtuple,

gamma: float,

nstep: int = 1,

cum_reward: bool = False,

value_gamma: Optional[torch.Tensor] = None,

criterion: torch.nn.modules = nn.MSELoss(reduction='none'),

) -> torch.Tensor:

"""

q, next_n_q, action, next_n_action, reward, done, weight = data

assert len(action.shape) == 1, action.shape

if weight is None:

weight = torch.ones_like(action)

batch_range = torch.arange(action.shape[0])

q_s_a = q[batch_range, action]

target_q_s_a = next_n_q[batch_range, next_n_action]

if cum_reward:

if value_gamma is None:

target_q_s_a = reward + (gamma ** nstep) * target_q_s_a * (1 - done)

else:

target_q_s_a = reward + value_gamma * target_q_s_a * (1 - done)

else:

target_q_s_a = nstep_return(nstep_return_data(reward, target_q_s_a, done), gamma, nstep, value_gamma)

td_error_per_sample = criterion(q_s_a, target_q_s_a.detach())

data: namedtuple,

gamma: float,

lambda1: tuple,

lambda2: tuple,

margin_function: float,

nstep: int = 1,

cum_reward: bool = False,

value_gamma: Optional[torch.Tensor] = None,

criterion: torch.nn.modules = nn.MSELoss(reduction='none'),

) -> torch.Tensor:

"""

q, next_n_q, action, next_n_action, reward, done, done_1, weight, new_n_q_one_step, next_n_action_one_step,\

is_expert = data # set is_expert flag(expert 1, agent 0)

assert len(action.shape) == 1, action.shape

if weight is None:

weight = torch.ones_like(action)

batch_range = torch.arange(action.shape[0])

q_s_a = q[batch_range, action]

target_q_s_a = next_n_q[batch_range, next_n_action]

target_q_s_a_one_step = new_n_q_one_step[batch_range, next_n_action_one_step]

# calculate n-step TD-loss

if cum_reward:

if value_gamma is None:

target_q_s_a = reward + (gamma ** nstep) * target_q_s_a * (1 - done)

else:

target_q_s_a = reward + value_gamma * target_q_s_a * (1 - done)

else:

target_q_s_a = nstep_return(nstep_return_data(reward, target_q_s_a, done), gamma, nstep, value_gamma)

td_error_per_sample = criterion(q_s_a, target_q_s_a.detach())

# calculate 1-step TD-loss

nstep = 1

reward = reward[0].unsqueeze(0)

value_gamma = None

if cum_reward:

if value_gamma is None:

target_q_s_a_one_step = reward + (gamma ** nstep) * target_q_s_a_one_step * (1 - done_1)

else:

target_q_s_a_one_step = reward + value_gamma * target_q_s_a_one_step * (1 - done_1)

else:

target_q_s_a_one_step = nstep_return(

nstep_return_data(reward, target_q_s_a_one_step, done_1), gamma, nstep, value_gamma

)

td_error_one_step_per_sample = criterion(q_s_a, target_q_s_a_one_step.detach())

# calculate the supervised loss

device = q_s_a.device

device_cpu = torch.device('cpu')

'''

l = margin_function * torch.ones_like(q).to(device_cpu)

l.scatter_(

1, torch.LongTensor(action.unsqueeze(1).to(device_cpu)), torch.zeros_like(q, device=device_cpu)

) # along the first dimension. for the index of the action, fill the corresponding position in l with 0

JE = is_expert * (torch.max(q + l.to(device), dim=1)[0] - q_s_a)

'''

return ((lambda1[0] * td_error_per_sample + td_error_one_step_per_sample + lambda2[0] * JE) *

r"""

if len(args) <= 0:

tmp = [kwargs['data'].reward.shape[0]]

tmp.extend(list(kwargs['data'].q.shape))

else:

tmp = [args[0].reward.shape[0]]

tmp.extend(list(args[0].q.shape))

data: namedtuple,

gamma: float,

nstep: int = 1,

value_gamma: Optional[torch.Tensor] = None,

criterion: torch.nn.modules = nn.MSELoss(reduction='none'),

trans_fn: Callable = value_transform,

inv_trans_fn: Callable = value_inv_transform,

) -> torch.Tensor:

"""

q, next_n_q, action, next_n_action, reward, done, weight = data

assert len(action.shape) == 1, action.shape

if weight is None:

weight = torch.ones_like(action)

batch_range = torch.arange(action.shape[0])

q_s_a = q[batch_range, action]

target_q_s_a = next_n_q[batch_range, next_n_action]

target_q_s_a = inv_trans_fn(target_q_s_a)

target_q_s_a = nstep_return(nstep_return_data(reward, target_q_s_a, done), gamma, nstep, value_gamma)

target_q_s_a = trans_fn(target_q_s_a)

td_error_per_sample = criterion(q_s_a, target_q_s_a.detach())

data: namedtuple,

gamma: float,

nstep: int = 1,

value_gamma: Optional[torch.Tensor] = None,

) -> torch.Tensor:

"""

q, next_n_q, action, next_n_action, reward, done, tau, weight = data

assert len(action.shape) == 1, action.shape

assert len(next_n_action.shape) == 1, next_n_action.shape

assert len(done.shape) == 1, done.shape

assert len(q.shape) == 3, q.shape

assert len(next_n_q.shape) == 3, next_n_q.shape

assert len(reward.shape) == 2, reward.shape

if weight is None:

weight = torch.ones_like(action)

batch_range = torch.arange(action.shape[0])

# shape: batch_size x num x 1

q_s_a = q[batch_range, action, :].unsqueeze(2)

# shape: batch_size x 1 x num

target_q_s_a = next_n_q[batch_range, next_n_action, :].unsqueeze(1)

assert reward.shape[0] == nstep

reward_factor = torch.ones(nstep).to(reward)

for i in range(1, nstep):

reward_factor[i] = gamma * reward_factor[i - 1]

# shape: batch_size

reward = torch.matmul(reward_factor, reward)

# shape: batch_size x 1 x num

if value_gamma is None:

target_q_s_a = reward.unsqueeze(-1).unsqueeze(-1) + (gamma ** nstep

) * target_q_s_a * (1 - done).unsqueeze(-1).unsqueeze(-1)

else:

target_q_s_a = reward.unsqueeze(-1).unsqueeze(

-1

) + value_gamma.unsqueeze(-1).unsqueeze(-1) * target_q_s_a * (1 - done).unsqueeze(-1).unsqueeze(-1)

# shape: batch_size x num x num

u = F.smooth_l1_loss(target_q_s_a, q_s_a, reduction="none")

# shape: batch_size

loss = (u * (tau - (target_q_s_a - q_s_a).detach().le(0.).float()).abs()).sum(-1).mean(1)

data: namedtuple,

gamma: float,

alpha: float,

nstep: int = 1,

cum_reward: bool = False,

value_gamma: Optional[torch.Tensor] = None,

criterion: torch.nn.modules = nn.MSELoss(reduction='none'),

) -> torch.Tensor:

"""

q, next_n_q, action, next_n_action, reward, done, weight = data

assert len(action.shape) == 1, action.shape

if weight is None:

weight = torch.ones_like(action)

batch_range = torch.arange(action.shape[0])

q_s_a = q[batch_range, action]

# target_q_s_a = next_n_q[batch_range, next_n_action]

target_v = alpha * torch.logsumexp(

next_n_q / alpha, 1

) # target_v = alpha * torch.log(torch.sum(torch.exp(next_n_q / alpha), 1))

target_v[target_v == float("Inf")] = 20

target_v[target_v == float("-Inf")] = -20

# For an appropriate hyper-parameter alpha, these hardcodes can be removed.

# However, algorithms may face the danger of explosion for other alphas.

# The hardcodes above are to prevent this situation from happening

record_target_v = copy.deepcopy(target_v)

#print(target_v)

if cum_reward:

if value_gamma is None:

target_v = reward + (gamma ** nstep) * target_v * (1 - done)

else:

target_v = reward + value_gamma * target_v * (1 - done)

else:

target_v = nstep_return(nstep_return_data(reward, target_v, done), gamma, nstep, value_gamma)

td_error_per_sample = criterion(q_s_a, target_v.detach())

data: namedtuple,

gamma: float,

nstep: int = 1,

kappa: float = 1.0,

value_gamma: Optional[torch.Tensor] = None,

) -> torch.Tensor:

"""

q, next_n_q, action, next_n_action, reward, done, replay_quantiles, weight = data

assert len(action.shape) == 1, action.shape

assert len(next_n_action.shape) == 1, next_n_action.shape

assert len(done.shape) == 1, done.shape

assert len(q.shape) == 3, q.shape

assert len(next_n_q.shape) == 3, next_n_q.shape

assert len(reward.shape) == 2, reward.shape

if weight is None:

weight = torch.ones_like(action)

batch_size = done.shape[0]

tau = q.shape[0]

tau_prime = next_n_q.shape[0]

action = action.repeat([tau, 1]).unsqueeze(-1)

next_n_action = next_n_action.repeat([tau_prime, 1]).unsqueeze(-1)

# shape: batch_size x tau x a

q_s_a = torch.gather(q, -1, action).permute([1, 0, 2])

# shape: batch_size x tau_prim x 1

target_q_s_a = torch.gather(next_n_q, -1, next_n_action).permute([1, 0, 2])

assert reward.shape[0] == nstep

device = torch.device("cuda" if reward.is_cuda else "cpu")

reward_factor = torch.ones(nstep).to(device)

for i in range(1, nstep):

reward_factor[i] = gamma * reward_factor[i - 1]

reward = torch.matmul(reward_factor, reward)

if value_gamma is None:

target_q_s_a = reward.unsqueeze(-1) + (gamma ** nstep) * target_q_s_a.squeeze(-1) * (1 - done).unsqueeze(-1)

else:

target_q_s_a = reward.unsqueeze(-1) + value_gamma.unsqueeze(-1) * target_q_s_a.squeeze(-1) * (1 - done

).unsqueeze(-1)

target_q_s_a = target_q_s_a.unsqueeze(-1)

# shape: batch_size x tau' x tau x 1.

bellman_errors = (target_q_s_a[:, :, None, :] - q_s_a[:, None, :, :])

# The huber loss (see Section 2.3 of the paper) is defined via two cases:

huber_loss = torch.where(

bellman_errors.abs() <= kappa, 0.5 * bellman_errors ** 2, kappa * (bellman_errors.abs() - 0.5 * kappa)

)

# Reshape replay_quantiles to batch_size x num_tau_samples x 1

replay_quantiles = replay_quantiles.reshape([tau, batch_size, 1]).permute([1, 0, 2])

# shape: batch_size x num_tau_prime_samples x num_tau_samples x 1.

replay_quantiles = replay_quantiles[:, None, :, :].repeat([1, tau_prime, 1, 1])

# shape: batch_size x tau_prime x tau x 1.

quantile_huber_loss = (torch.abs(replay_quantiles - ((bellman_errors < 0).float()).detach()) * huber_loss) / kappa

# shape: batch_size

loss = quantile_huber_loss.sum(dim=2).mean(dim=1)[:, 0]

r"""

if len(args) <= 0:

tmp = kwargs['data'].reward.shape[0]

else:

tmp = args[0].reward.shape

"""

value, reward, weight = data

if weight is None:

weight = torch.ones_like(reward)

with torch.no_grad():

return_ = generalized_lambda_returns(value, reward, gamma, lambda_)

# discard the value at T as it should be considered in the next slice

loss = 0.5 * \

(F.mse_loss(return_, value[:-1], reduction='none') * weight).mean()

bootstrap_values: torch.Tensor, rewards: torch.Tensor, gammas: float, lambda_: float

) -> torch.Tensor:

r"""

if not isinstance(gammas, torch.Tensor):

gammas = gammas * torch.ones_like(rewards)

if not isinstance(lambda_, torch.Tensor):

lambda_ = lambda_ * torch.ones_like(rewards)

bootstrap_values_tp1 = bootstrap_values[1:, :]

bootstrap_values: torch.Tensor, rewards: torch.Tensor, gammas: float, lambda_: float

) -> torch.Tensor:

r"""

result = torch.empty_like(rewards)

# Forced cutoff at the last one

result[-1, :] = rewards[-1, :] + gammas[-1, :] * bootstrap_values[-1, :]

discounts = gammas * lambda_

for t in reversed(range(rewards.size()[0] - 1)):

result[t, :] = rewards[t, :] \

+ discounts[t, :] * result[t + 1, :] \

+ (gammas[t, :] - discounts[t, :]) * bootstrap_values[t, :]