Spaces:
Running
Running
| from __future__ import absolute_import | |
| from __future__ import division | |
| from __future__ import print_function | |
| """Manage data for pretraining and RL tasks.""" | |
| import ast | |
| from collections import namedtuple | |
| from absl import logging | |
| from single_task import code_tasks # brain coder | |
| RLBatch = namedtuple('RLBatch', ['reward_fns', 'batch_size', 'good_reward']) | |
| class DataManager(object): | |
| """Interface between environment and model.""" | |
| def __init__(self, global_config, run_number=None, | |
| do_code_simplification=False): | |
| """Constructs a DataManager. | |
| Args: | |
| global_config: A config_lib.Config instance containing all config. See | |
| config in defaults.py. | |
| run_number: Which run this is (of the same experiment). This should be set | |
| when a task cycle is defined in the config. A task cycle is a list of | |
| tasks to cycle through repeatedly, and the selected task is a function | |
| of the run number, i.e. 0-th run, 1-st run, 2-nd run, etc... | |
| This can be None if only a single task is set in the config. | |
| do_code_simplification: When global_config.env.config_for_iclr is True, | |
| use this option to create code simplification (code golf) tasks, vs | |
| fixed length coding tasks. If True, a task with code simplification | |
| reward will be constructed. | |
| Raises: | |
| ValueError: If global_config.env.task and global_config.env.task_cycle | |
| are both set, or both not set. Only one should be given. | |
| ValueError: If global_config.env.task_cycle is set but run_number is None. | |
| """ | |
| env_config = global_config.env | |
| self.batch_size = global_config.batch_size | |
| if env_config.task_cycle: | |
| if env_config.task: | |
| raise ValueError('Do not set both `task` and `task_cycle`.') | |
| if run_number is None: | |
| raise ValueError('Do not use task_cycle for single-run experiment.') | |
| index = run_number % len(env_config.task_cycle) | |
| self.task_name = env_config.task_cycle[index] | |
| logging.info('run_number: %d, task_cycle index: %d', run_number, index) | |
| logging.info('task_cycle: %s', env_config.task_cycle) | |
| elif env_config.task: | |
| self.task_name = env_config.task | |
| else: | |
| raise ValueError('Either `task` or `task_cycle` must be set.') | |
| logging.info('Task for this run: "%s"', self.task_name) | |
| logging.info('config_for_iclr=True; do_code_simplification=%s', | |
| do_code_simplification) | |
| self.rl_task = code_tasks.make_task( | |
| task_name=self.task_name, | |
| override_kwargs=ast.literal_eval(env_config.task_kwargs), | |
| max_code_length=global_config.timestep_limit, | |
| require_correct_syntax=env_config.correct_syntax, | |
| do_code_simplification=do_code_simplification, | |
| correct_bonus=env_config.task_manager_config.correct_bonus, | |
| code_length_bonus=env_config.task_manager_config.code_length_bonus) | |
| def sample_rl_batch(self): | |
| """Create reward functions from the current task. | |
| Returns: | |
| RLBatch namedtuple instance, which holds functions and information for | |
| a minibatch of episodes. | |
| * reward_fns: A reward function for each episode. Maps code string to | |
| reward. | |
| * batch_size: Number of episodes in this minibatch. | |
| * good_reward: Estimated threshold of rewards which indicate the algorithm | |
| is starting to solve the task. This is a heuristic that tries to | |
| reduce the amount of stuff written to disk. | |
| """ | |
| reward_fns = self.rl_task.rl_batch(self.batch_size) | |
| return RLBatch( | |
| reward_fns=reward_fns, | |
| batch_size=self.batch_size, | |
| good_reward=self.rl_task.good_reward) | |