File size: 1,514 Bytes
ce3feed |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
import torch
def log_metrics(metrics, step, tb_writer, wandb_writer):
# Log metrics to TensorBoard
if tb_writer:
for key, value in metrics.items():
tb_writer.add_scalar(key, value, step)
# Log metrics to wandb
# if wandb_writer:
# wandb_writer.log(metrics, step=step)
def td_lambda(rewards, predicted_discount, values, lambda_, device):
"""
Compute the TD(位) returns for value estimation.
Args:
- rewards (Tensor): Tensor of rewards with shape [batch_size, horizon_len, 1].
- predicted_discount (Tensor): Tensor indicating probability of episode termination with shape [batch_size, horizon_len, 1].
- values (Tensor): Tensor of value estimates with shape [batch_size, horizon_len, 1].
- lambda_ (float): The 位 parameter in TD(位) controlling bias-variance tradeoff.
Returns:
- td_lambda (Tensor): The computed lambda returns with shape [batch_size, time_steps - 1].
"""
batch_size, _, _ = rewards.shape
last_lambda = torch.zeros((batch_size, 1)).to(device)
cur_rewards = rewards[:, :-1]
next_values = values[:, 1:]
predicted_discount = predicted_discount[:, :-1]
td_1 = cur_rewards + predicted_discount * next_values * (1 - lambda_)
returns = torch.zeros_like(cur_rewards).to(device)
for i in reversed(range(td_1.size(1))):
last_lambda = td_1[:, i] + predicted_discount[:, i] * lambda_ * last_lambda
returns[:, i] = last_lambda
return returns |