File size: 3,685 Bytes
97b6013 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
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)
|