check for converted data as specified by config and generate it if necessary (#9)

Author: juliusfrost

envs/data.py

@@ -54,7 +54,7 @@ def wrap_env(env: MinerRLDataEnv, env_config):
     return wrap(env, **env_config)
 
 
-def write_jsons(environment, data_dir, env_config, save_path, overwrite=False, **kwargs):
+def write_jsons(environment, data_dir, env_config, save_path, overwrite=False, fail_safe=True, **kwargs):
     data_pipeline = minerl.data.make(environment, data_dir, **kwargs)
     env = MinerRLDataEnv(data_pipeline)
     env = wrap_env(env, env_config)
@@ -67,8 +67,12 @@ def write_jsons(environment, data_dir, env_config, save_path, overwrite=False, *
             print(f'Overwriting! {abs_save_path}')
             shutil.rmtree(abs_save_path)
         else:
-            raise ValueError(f'Directory {abs_save_path} not empty!'
-                             f'Cannot overwrite existing data automatically, please delete old data if unused.')
+            if fail_safe:
+                print(f'Json data already exists at {abs_save_path}')
+                return
+            else:
+                raise ValueError(f'Directory {abs_save_path} not empty!'
+                                 f'Cannot overwrite existing data automatically, please delete old data if unused.')
 
     batch_builder = SampleBatchBuilder()
     writer = JsonWriter(save_path)

generate_kmeans.py

@@ -11,6 +11,7 @@
 parser.add_argument('--env', default=None)
 parser.add_argument('--num-actions', type=int, default=32)
 parser.add_argument('--data-dir', default=os.getenv('MINERL_DATA_ROOT', 'data'))
+parser.add_argument('--overwrite', action='store_true')
 
 
 def main():
@@ -25,6 +26,15 @@ def main():
 
     for env_name in env_list:
         print(f'Generating {args.num_actions}-means for {env_name}')
+
+        file_dir = os.path.join(args.data_dir, f'{args.num_actions}-means')
+        file = os.path.join(file_dir, env_name + '.npy')
+        if os.path.exists(file) and not args.overwrite:
+            print(f'k-means file already exists at {file}')
+            continue
+        if not os.path.exists(file_dir):
+            os.mkdir(file_dir)
+
         data = minerl.data.make(env_name)
         actions = []
         for trajectory_name in tqdm(list(data.get_trajectory_names())):
@@ -37,10 +47,6 @@ def main():
         print('computing k-means...')
         kmeans = KMeans(n_clusters=args.num_actions, verbose=1, random_state=0).fit(actions)
         print(kmeans)
-        file_dir = os.path.join(args.data_dir, f'{args.num_actions}-means')
-        if not os.path.exists(file_dir):
-            os.mkdir(file_dir)
-        file = os.path.join(file_dir, env_name + '.npy')
         np.save(file, kmeans.cluster_centers_)
 
 

rllib_train.py

@@ -1,11 +1,11 @@
 """
 Runs the training process through the RLLib train script
-Example command:
-$ python rllib_train.py -f config/minerl-impala-debug.yaml
 See help with:
 $ python rllib_train.py --help
 """
+import subprocess
 
+import yaml
 from ray.rllib.train import create_parser, run
 
 import envs
@@ -15,9 +15,39 @@
 models.register()
 
 
+def generate_kmeans(env):
+    command = f'python generate_kmeans.py --env {env}'
+    print('running:', command)
+    subprocess.run(command.split())
+
+
+def convert_data(args):
+    command = f'python convert_data.py -f {args.config_file}'
+    print('running:', command)
+    subprocess.run(command.split())
+
+
+def check_data(args):
+    if args.config_file is not None:
+        config = yaml.safe_load(open(args.config_file))
+        settings = list(config.values())[0]
+        if 'env' in settings:
+            env = settings['env']
+        if 'config' in settings:
+            if 'env' in settings['config']:
+                env = settings['config']['env']
+            if 'env_config' in settings['config']:
+                env_config = settings['config']['env_config']
+                if env_config.get('discrete', False):
+                    generate_kmeans(env)
+            if 'input' in settings['config']:
+                convert_data(args)
+
+
 def main():
     parser = create_parser()
     args = parser.parse_args()
+    check_data(args)
     run(args, parser)