| 14 | |
| 15 | @ray.remote(num_gpus=0.125) |
| 16 | class DataWorker(object): |
| 17 | def __init__(self, rank, replay_buffer, storage, config): |
| 18 | """Data Worker for collecting data through self-play |
| 19 | Parameters |
| 20 | ---------- |
| 21 | rank: int |
| 22 | id of the worker |
| 23 | replay_buffer: Any |
| 24 | Replay buffer |
| 25 | storage: Any |
| 26 | The model storage |
| 27 | """ |
| 28 | self.rank = rank |
| 29 | self.config = config |
| 30 | self.storage = storage |
| 31 | self.replay_buffer = replay_buffer |
| 32 | # double buffering when data is sufficient |
| 33 | self.trajectory_pool = [] |
| 34 | self.pool_size = 1 |
| 35 | self.device = self.config.device |
| 36 | self.gap_step = self.config.num_unroll_steps + self.config.td_steps |
| 37 | self.last_model_index = -1 |
| 38 | |
| 39 | def put(self, data): |
| 40 | # put a game history into the pool |
| 41 | self.trajectory_pool.append(data) |
| 42 | |
| 43 | def len_pool(self): |
| 44 | # current pool size |
| 45 | return len(self.trajectory_pool) |
| 46 | |
| 47 | def free(self): |
| 48 | # save the game histories and clear the pool |
| 49 | if self.len_pool() >= self.pool_size: |
| 50 | self.replay_buffer.save_pools.remote(self.trajectory_pool, self.gap_step) |
| 51 | del self.trajectory_pool[:] |
| 52 | |
| 53 | def put_last_trajectory(self, i, last_game_histories, last_game_priorities, game_histories): |
| 54 | """put the last game history into the pool if the current game is finished |
| 55 | Parameters |
| 56 | ---------- |
| 57 | last_game_histories: list |
| 58 | list of the last game histories |
| 59 | last_game_priorities: list |
| 60 | list of the last game priorities |
| 61 | game_histories: list |
| 62 | list of the current game histories |
| 63 | """ |
| 64 | # pad over last block trajectory |
| 65 | beg_index = self.config.stacked_observations |
| 66 | end_index = beg_index + self.config.num_unroll_steps |
| 67 | |
| 68 | pad_obs_lst = game_histories[i].obs_history[beg_index:end_index] |
| 69 | pad_child_visits_lst = game_histories[i].child_visits[beg_index:end_index] |
| 70 | |
| 71 | beg_index = 0 |
| 72 | end_index = beg_index + self.gap_step - 1 |
| 73 |
nothing calls this directly
no outgoing calls
no test coverage detected