code: def _adjust_nstep(n_step, gamma, obs, actions, rewards, new_obs, dones):
      
  
      assert not any(dones[:-1]), "Unexpected done in middle of trajectory"
  
      traj_length = len(rewards)
      for i in range(traj_length):
          for j in range(1, n_step):
              if i + j < traj_length:
                  new_obs[i] = new_obs[i + j]
                  dones[i] = dones[i + j]
                  rewards[i] += gamma**j * rewards[i + j]
docs: Rewrites the given trajectory fragments to encode n-step rewards.
  
      reward[i] = (
          reward[i] * gamma**0 +
          reward[i+1] * gamma**1 +
          ... +
          reward[i+n_step-1] * gamma**(n_step-1))
  
      The ith new_obs is also adjusted to point to the (i+n_step-1)'th new obs.
  
      At the end of the trajectory, n is truncated to fit in the traj length.