import torch


def compute_rtgs(batch_rews):
	"""
		Compute the Reward-To-Go of each timestep in a batch given the rewards.

		Parameters:
			batch_rews - the rewards in a batch, Shape: (number of episodes, number of timesteps per episode)

		Return:
			batch_rtgs - the rewards to go, Shape: (number of timesteps in batch)
	"""
	# The rewards-to-go (rtg) per episode per batch to return.
	# The shape will be (num timesteps per episode)
	batch_rtgs = []

	# Iterate through each episode
	for ep_rews in reversed(batch_rews):

		discounted_reward = 0 # The discounted reward so far

		# Iterate through all rewards in the episode. We go backwards for smoother calculation of each
		# discounted return (think about why it would be harder starting from the beginning)
		for rew in reversed(ep_rews):
			discounted_reward = rew + discounted_reward * 0.95
			batch_rtgs.insert(0, discounted_reward)

	# Convert the rewards-to-go into a tensor
	batch_rtgs = torch.tensor(batch_rtgs, dtype=torch.float)

	return batch_rtgs