def _compute_path_probs(paths, pol_dist_type=None, insert=True, insert_key='a_logprobs'): """ Returns a N x T matrix of action probabilities """ if pol_dist_type is None: # try to infer distribution type path0 = paths[0] if 'log_std' in path0['agent_infos']: pol_dist_type = DIST_GAUSSIAN elif 'prob' in path0['agent_infos']: pol_dist_type = DIST_CATEGORICAL else: raise NotImplementedError() # compute path probs Npath = len(paths) actions = [path['actions'] for path in paths] if pol_dist_type == DIST_GAUSSIAN: params = [(path['agent_infos']['mean'], path['agent_infos']['log_std']) for path in paths] path_probs = [gauss_log_pdf(params[i], actions[i]) for i in range(Npath)] elif pol_dist_type == DIST_CATEGORICAL: params = [(path['agent_infos']['prob'],) for path in paths] path_probs = [categorical_log_pdf(params[i], actions[i]) for i in range(Npath)] else: raise NotImplementedError("Unknown distribution type") if insert: for i, path in enumerate(paths): path[insert_key] = path_probs[i] return np.array(path_probs)