def get_dataset_params(folders):
# NOTE: you can pass in arguments so that you can reuse,
# such as using this function to create the training and holdout params
return d(
cls=Dataset,
params=d(
# TODO
))
def get_env_spec_params():
# NOTE: each "sub" param has a class and a params that's passed into it
# this separation is because the config file should not instantiate any classes, only define
return d(
cls=EnvSpec,
params=d(
# TODO
))
def get_trainer_params():
return d(
cls=Trainer,
params=d(
max_steps=int(1e5),
step_train_env_every_n_steps=1,
step_holdout_env_every_n_steps=50,
holdout_every_n_steps=50,
log_every_n_steps=int(1e3),
save_every_n_steps=int(1e4),
cost_fn=lambda model_outputs, outputs: d(),
# TODO
))
def get_params():
return d(
exp_name='bumpy',
# NOTE: this is where all the params get created
env_spec=get_env_spec_params(),
env=get_env_params(),
dataset_train=get_dataset_params(['/some/training/folders']),
dataset_holdout=get_dataset_params(['/some/holdout/folders']),
model=get_model_params(),
trainer=get_trainer_params(),
policy=get_policy_params(),
)
if self._data_len - self._last_save >= self._save_every_n_steps:
print("SAVING:", self._data_len)
self.save()
self._last_save = self._data_len
# don't call this function too often
def save(self):
save_dict = {}
for key in self._all_names:
save_dict[key] = np.concatenate(self._datadict[key])
path = os.path.join(self._file_manager.exp_dir, self._output_file)
np.savez_compressed(path, **save_dict)
def __len__(self):
return self._data_len
if __name__ == '__main__':
from mbrl.envs.env_spec import EnvSpec
from dotmap import DotMap as d
d = NpDataset(
d(
input_file=None,
output_file='delete',
batch_size=5,
horizon=9, # effective batch of batch_size*(horizon+1)
),
EnvSpec())
def get_policy_params():
return d(
cls=Policy,
params=d(
# TODO
))
def get_model_params():
return d(
cls=Model,
params=d(
# TODO
))
def get_env_params():
return d(
cls=Env,
params=d(
# TODO
))