import torch


def calculate_R_pi(R, pi):
    """
    calculates R_pi
    R_pi(s) = \sum_a pi(s,a) r(s,a)
    :param R: reward matrix of size |S| x |A|
    :param pi: matrix of size |S| x |A| indicating the policy
    :return:
    """
    return torch.einsum('sa,sa->s', R, pi)