def td0_estimate(rewards, next_states, dones, V, gamma=1.):
    """
    TD(0) estimate of current state's value given observations of reward and
    next state.

    Works on single observations or batches.

    Params
    ======
    rewards (Tensor): reward(s) for taking current action(s)
    next_states (Tensor): next state(s) after taking current action(s)
    V (nn.Module): state value function
    gamma (float): discount factor
    """
    return rewards + gamma * (1. - dones) * V(next_states)