import torch


def expected_sarsa_estimate(rewards, next_states, dones, Q, policy, gamma=1.):
    """
    Expected Sarsa estimator of the value of the current state-action pair.
    See docstring for `sarsa_estimate` for further details.

    The policy is expected to be a StochasticPolicy that generates Categorical
    distributions when the `stoch_act` method is called.
    """
    pi = policy.stoch_act(next_states).probs
    expected_next_values = torch.sum(pi * Q(next_states), dim=-1)
    return rewards + gamma * dones * expected_next_values