import torch


def calculate_P_pi(P, pi):
    """
    calculates P_pi
    P_pi(s,t) = \sum_a pi(s,a) p(s, a, t)
    :param P: transition matrix of size |S|x|A|x|S|
    :param pi: matrix of size |S| x |A| indicating the policy
    :return: a matrix of size |S| x |S|
    """
    return torch.einsum('sat,sa->st', P, pi)