fix dreamerv2 config

dreamerv2/core/config.py

@@ -32,7 +32,8 @@ def save(self, filename):
     elif filename.suffix in ('.yml', '.yaml'):
       import ruamel.yaml as yaml
       with io.StringIO() as stream:
-        yaml.safe_dump(dict(self), stream)
+        yaml = yaml.YAML(typ='safe', pure=True)
+        yaml.dump(dict(self), stream)
         filename.write(stream.getvalue())
     else:
       raise NotImplementedError(filename.suffix)
@@ -44,7 +45,8 @@ def load(cls, filename):
       return cls(json.loads(filename.read_text()))
     elif filename.suffix in ('.yml', '.yaml'):
       import ruamel.yaml as yaml
-      return cls(yaml.safe_load(filename.read_text()))
+      yaml = yaml.YAML(typ='safe', pure=True)
+      return cls(yaml.load(filename.read_text()))
     else:
       raise NotImplementedError(filename.suffix)
 
@@ -174,8 +176,8 @@ def _ensure_values(self, mapping):
         if len(value) == 0:
           message = 'Empty lists are disallowed because their type is unclear.'
           raise TypeError(message)
-        if not isinstance(value[0], (str, float, int, bool, list)):
-          message = 'Lists can only contain strings, floats, ints, bools, lists'
+        if not isinstance(value[0], (str, float, int, bool, list, dict)):
+          message = 'Lists can only contain strings, floats, ints, bools, lists, dict'
           message += f' but not {type(value[0])}'
           raise TypeError(message)
         if not all(isinstance(x, type(value[0])) for x in value[1:]):