def sarsa_estimate(rewards, next_states, dones, Q, policy, gamma=1.,
                   discrete_actions=True):
    """
    Sarsa estimator of the value of the current state-action pair given
    observations of the reward and next state.

    Can work on single observations or batches.

    Note this function expects the `policy` to be consistent with Q (i.e. it
    won't check this). True Sarsa requires `policy` to be the policy currently
    used by the agent, but this function can also be used for deep Q-learning
    or DDPG, i.e. off-policy.

    If `discrete_actions` is true, the Q function is assumed to map states
    to action value arrays; if `discrete_actions` is false, the Q function
    is assumed to map state-action value pairs to corresponding value estimates.

    Params
    ======
    rewards (Tensor): reward(s) for taking current action(s)
    next_states (Tensor): state(s) reached after taking current action(s)
    dones (Tensor): done flag(s) for next state(s)
    Q (nn.Module): Q module with properties described above
    policy (nn.Module): policy derived from Q (e.g. `EpsilonGreedyPolicy`)
    gamma (float): discount factor
    discrete_actions (bool): determines the action space type    
    """
    
    next_actions = policy.act(next_states)
    if discrete_actions:
        # The policy may have collapsed dimensions by one (i.e. passing a
        # single state returns a scalar action, passing an array of states
        # returns a 1D array), so we may need to unsqueeze to prepare for gather
        if len(next_actions.shape) < len(next_states.shape):
            next_actions = next_actions.unsqueeze(-1)
        next_values = Q(next_states).gather(-1, next_actions)
        next_values = next_values.squeeze() # collapse dimensions again
    else:
        next_values = Q(next_states, next_actions).squeeze()
    return rewards + gamma * (1. - dones) * next_values