Provides a common generate_examples method for D4RL datasets.
(file_path: str)
| 53 | |
| 54 | |
| 55 | def generate_examples(file_path: str): |
| 56 | """Provides a common generate_examples method for D4RL datasets.""" |
| 57 | d4rl_dict = read_d4rl_dataset(file_path) |
| 58 | if 'timeouts' not in d4rl_dict: |
| 59 | raise ValueError('Only datasets with explicit timeouts are supported.') |
| 60 | |
| 61 | done = [ |
| 62 | terminal or timeout |
| 63 | for (terminal, timeout) in zip( |
| 64 | d4rl_dict['terminals'], d4rl_dict['timeouts'] |
| 65 | ) |
| 66 | ] |
| 67 | # is_first corresponds to the done flag delayed by one step. |
| 68 | d4rl_dict['is_first'] = [True] + done[:-1] |
| 69 | # is_last is not used but this is needed to build a valid dictionary. |
| 70 | d4rl_dict['is_last'] = done |
| 71 | |
| 72 | # Get step metadata |
| 73 | infos_dict = _get_nested_metadata(d4rl_dict, 'infos') |
| 74 | |
| 75 | # Flatten reward |
| 76 | d4rl_dict['rewards'] = np.squeeze(d4rl_dict['rewards']) |
| 77 | |
| 78 | episode_metadata = _get_nested_metadata(d4rl_dict, 'metadata') |
| 79 | dataset_dict = { |
| 80 | 'observation': d4rl_dict['observations'], |
| 81 | 'action': d4rl_dict['actions'], |
| 82 | 'reward': d4rl_dict['rewards'], |
| 83 | 'discount': np.ones_like(d4rl_dict['rewards']), |
| 84 | 'is_terminal': d4rl_dict['terminals'], |
| 85 | 'is_first': d4rl_dict['is_first'], |
| 86 | 'is_last': d4rl_dict['is_last'], |
| 87 | } |
| 88 | if 'next_observations' in d4rl_dict: |
| 89 | dataset_dict['next_observation'] = d4rl_dict['next_observations'] |
| 90 | |
| 91 | if infos_dict: |
| 92 | dataset_dict['infos'] = infos_dict |
| 93 | num_steps = len(dataset_dict['is_first']) |
| 94 | prev = 0 |
| 95 | counter = 0 |
| 96 | for pos in range(num_steps): |
| 97 | if dataset_dict['is_first'][pos] and pos > prev: |
| 98 | yield counter, _get_episode(dataset_dict, episode_metadata, prev, pos) |
| 99 | prev = pos |
| 100 | counter += 1 |
| 101 | if prev < num_steps: |
| 102 | yield counter, _get_episode(dataset_dict, episode_metadata, prev, num_steps) |
| 103 | |
| 104 | |
| 105 | def _get_nested_metadata( |
nothing calls this directly
no test coverage detected