Replace lambda with partial in env instantiation.

Following hexaflexa suggestion in Issue #4.

src/play.py

@@ -1,3 +1,4 @@
+from functools import partial 
 from pathlib import Path
 
 import hydra
@@ -18,7 +19,7 @@ def main(cfg: DictConfig):
     assert cfg.mode in ('world_model', 'episode_replay', 'agent')
 
     if cfg.mode in ['world_model', 'agent']:
-        env_fn = lambda: instantiate(cfg.env.test)
+        env_fn = partial(instantiate, config=cfg.env.test)
         test_env = SingleProcessEnv(env_fn)
         tokenizer = instantiate(cfg.tokenizer)
         world_model = WorldModel(obs_vocab_size=tokenizer.vocab_size, act_vocab_size=test_env.num_actions, config=instantiate(cfg.world_model))

src/trainer.py

@@ -1,4 +1,5 @@
 from collections import defaultdict
+from functools import partial
 from pathlib import Path
 import shutil
 import sys
@@ -62,7 +63,7 @@ class Trainer:
         self.episode_manager_imagination = EpisodeDirManager(self.episode_dir / 'imagination', max_num_episodes=cfg.evaluation.actor_critic.num_episodes_to_save)
 
         def create_env(cfg_env, num_envs):
-            env_fn = lambda: instantiate(cfg_env)
+            env_fn = partial(instantiate, config=cfg_env)
             return MultiProcessEnv(env_fn, num_envs, should_wait_num_envs_ratio=1.0) if num_envs > 1 else SingleProcessEnv(env_fn)
 
         if self.cfg.training.should: