to_n_step_transition returns wrong results if episode was truncated by the time limit wrapper and N > 1

[zhezherun]

If an episode is truncated by the time limit wrapper, the last discount in that episode is set to 1.0 instead of 0.0. As a result, both the reward and discount calculation spill over into the next episode and give incorrect values. The next observation is also taken from the end of the trajectory, but it should come from the end of the episode instead.

Please run the following code to reproduce the issue. In this code, both "short" and "long" trajectories should give exactly the same result because the episode was truncated.

import tensorflow as tf
import gym
from gym.wrappers import time_limit
import numpy as np
from tf_agents.environments import suite_gym, tf_py_environment
from tf_agents.drivers import tf_driver
from tf_agents.networks import network
from tf_agents.policies import actor_policy
from tf_agents.replay_buffers import tf_uniform_replay_buffer
from tf_agents.specs import tensor_spec
from tf_agents.trajectories import time_step as ts, policy_step, trajectory
from tf_agents.utils import nest_utils


class SimpleEnv(gym.Env):
    def __init__(self):
        self.action_space = gym.spaces.Box(
            low=np.array([-1.0], np.float32), high=np.array([1.0], np.float32), dtype=np.float32,
        )
        self.observation_space = gym.spaces.Box(
            low=np.array([0.0], np.float32), high=np.array([10.0], np.float32), dtype=np.float32,
        )
        self.idx = 0

    def reset(self):
        self.idx = 0
        return np.array([self.idx], np.float32)

    def step(self, action):
        self.idx += 1
        return np.array([self.idx], np.float32), 1.0, False, {}


class SimpleActor(network.Network):
    def call(self, observations, step_type=(), network_state=()):
        outer_shape = nest_utils.get_outer_shape(observations, self.input_tensor_spec)
        return tf.fill(tf.concat([outer_shape, [1]], axis=0), 0.0), network_state


def main():
    env = SimpleEnv()
    # Truncate episodes to 2 steps
    tf_env = tf_py_environment.TFPyEnvironment(suite_gym.wrap_env(env, max_episode_steps=2))

    policy_step_spec = policy_step.PolicyStep(action=tf_env.action_spec(), state=(), info=())
    time_step_spec = ts.time_step_spec(tf_env.observation_spec())
    trajectory_spec = trajectory.from_transition(time_step_spec, policy_step_spec, time_step_spec)
    replay_buffer = tf_uniform_replay_buffer.TFUniformReplayBuffer(
        trajectory_spec, batch_size=1, max_length=10
    )

    actor_network = SimpleActor(input_tensor_spec=tf_env.observation_spec())
    policy = actor_policy.ActorPolicy(
        time_step_spec=time_step_spec, action_spec=tf_env.action_spec(), actor_network=actor_network
    )
    driver = tf_driver.TFDriver(
        env=tf_env, policy=policy, observers=[replay_buffer.add_batch], max_steps=4
    )
    time_step = tf_env.reset()
    driver.run(time_step)

    assert replay_buffer.num_frames() == 5  # Including a boundary
    dataset = replay_buffer.as_dataset(sample_batch_size=1, num_steps=5)  # So that sampling is deterministic
    iterator = iter(dataset)

    experience, _buffer_info = next(iterator)
    print("Trajectory:", experience)  # Observations will be [0.0, 1.0, 2.0, 0.0, 1.0]

    short_experience = tf.nest.map_structure(lambda t: t[:, :3], experience)
    transition = trajectory.to_n_step_transition(short_experience, gamma=0.5)
    print("Short experience")
    print("Next observation", transition.next_time_step.observation)  # Expected observation: 2.0
    print("Discount:", transition.next_time_step.discount)  # Expected discount: 0.5
    print("Reward:", transition.next_time_step.reward)  # Expected reward: 1.0 + 0.5*1.0 = 1.5

    # Longer experience should give the same n_step transition because the episode is truncated
    transition = trajectory.to_n_step_transition(experience, gamma=0.5)
    print("Long experience")
    print("Next observation:", transition.next_time_step.observation)
    print("Discount", transition.next_time_step.discount)
    print("Reward:", transition.next_time_step.reward)

    tf_env.close()


if __name__ == "__main__":
    main()