Fix typing issue in base system.

og_marl/environments/base.py

@@ -53,7 +53,7 @@ def get_stats(self) -> Dict:
         """
         return {}
 
-    def render(self):
+    def render(self) -> Any:
         """Return frame for rendering"""
         return np.zeros((10, 10, 3), "float32")
 

og_marl/environments/jumanji_rware.py

@@ -119,7 +119,6 @@ def step(self, actions: Dict[str, np.ndarray]) -> StepReturn:
 
         return observations, rewards, terminals, truncations, info
 
-
     def render(self) -> Any:
         frame = self._state
 

og_marl/tf2/systems/base.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import time
-from typing import Dict, Optional
+from typing import Dict, Optional, Tuple, List
 
 import numpy as np
 from chex import Numeric
@@ -42,7 +42,7 @@ def __init__(
     def get_stats(self) -> Dict[str, Numeric]:
         return {}
 
-    def evaluate(self, num_eval_episodes: int = 4) -> Dict[str, Numeric]:
+    def evaluate(self, num_eval_episodes: int = 4) -> Tuple[Dict[str, Numeric], list]:
         """Method to evaluate the system online (i.e. in the environment)."""
         episode_returns = []
         all_frames = []
@@ -187,7 +187,7 @@ def train_offline(
         evaluate_every: int = 1000,
         num_eval_episodes: int = 4,
         json_writer: Optional[JsonWriter] = None,
-    ) -> None:
+    ) -> List:
         """Method to train the system offline.
 
         WARNING: make sure evaluate_every % log_every == 0 and log_every < evaluate_every,