How to do it...

Creating an MDP can be done via the following steps:

1. Import PyTorch and define the transition matrix:

 >>> import torch
 >>> T = torch.tensor([[[0.8, 0.1, 0.1],
 ...                    [0.1, 0.6, 0.3]],
 ...                   [[0.7, 0.2, 0.1],
 ...                    [0.1, 0.8, 0.1]],
 ...                   [[0.6, 0.2, 0.2],
 ...                    [0.1, 0.4, 0.5]]]
 ...                  )

1. Define the reward function and the discount factor:

 >>> R = torch.tensor([1., 0, -1.])
 >>> gamma = 0.5

1. The optimal policy in this case is selecting action a0 in all circumstances:

>>> action = 0

1. We calculate the value, V, of the optimal policy using the matrix inversion method in the following function:

 >>> def cal_value_matrix_inversion(gamma, trans_matrix, rewards):
 ...     inv = torch.inverse(torch.eye(rewards.shape[0]) 
                                             - gamma * trans_matrix)
 ...     V = torch.mm(inv, rewards.reshape(-1, 1))
 ...     return V

We will demonstrate how to derive the value in the next section.

1. We feed all variables we have to the function, including the transition probabilities associated with action a0:

 >>> trans_matrix = T[:, action]
 >>> V = cal_value_matrix_inversion(gamma, trans_matrix, R)
 >>> print("The value function under the optimal 
           policy is:\n{}".format(V))
 The value function under the optimal policy is:
 tensor([[ 1.6787],
         [ 0.6260],
         [-0.4820]])