code: def compute_probab_ratios(p_new, p_old, actions, reward_mask):
    
  
    B, T = actions.shape  # pylint: disable=invalid-name
    assert (B, T + 1) == p_old.shape[:2]
    assert (B, T + 1) == p_new.shape[:2]
  
    logp_old = chosen_probabs(p_old, actions)
    logp_new = chosen_probabs(p_new, actions)
  
    assert (B, T) == logp_old.shape
    assert (B, T) == logp_new.shape
  
    # Since these are log-probabilities, we just subtract them.
    probab_ratios = np.exp(logp_new - logp_old) * reward_mask
    assert (B, T) == probab_ratios.shape
    return probab_ratios
docs: Computes the probability ratios for each time-step in a trajectory.
  
    Args:
      p_new: ndarray of shape [B, T+1, A] of the log-probabilities that the policy
        network assigns to all the actions at each time-step in each batch using
        the old parameters.
      p_old: ndarray of shape [B, T+1, A], same as above, but using old policy
        network parameters.
      actions: ndarray of shape [B, T] where each element is from [0, A).
      reward_mask: ndarray of shape [B, T] masking over probabilities.
  
    Returns:
      probab_ratios: ndarray of shape [B, T], where
      probab_ratios_{b,t} = p_new_{b,t,action_{b,t}} / p_old_{b,t,action_{b,t}}