Trajectory visualization

.github/workflows/pythontests.yml

@@ -11,11 +11,12 @@ jobs:
         architecture: 'x64'
     - name: Install dependencies
       run: python -m pip install --upgrade pip
-    - name: Setup
-      run: pip install .
+    - name: Setup # needs -e flag to read overcooked_ai/visualization non python files
+      run: pip install -e .
     - name: Run tests
       run: |
-        python testing/overcooked_test.py 
+        python testing/overcooked_test.py
+        python testing/visualization_test.py
   ## This should be uncommented once the planners are updated ##
   # osx_full_tests:
   #   runs-on: macos-latest
@@ -46,11 +47,12 @@ jobs:
       run: |
         python -m pip install --upgrade pip
         pip install coverage
-    - name: Setup # needs -e flag for coverage tests to work
+    - name: Setup # needs -e flag for coverage tests to work and to read overcooked_ai/visualization non python files 
       run: pip install -e .
     - name: Run tests and generate coverage report
       run: |
         coverage run testing/overcooked_test.py
+        coverage run -a testing/visualization_test.py
         coverage report
     - name: Upload coverage to Codecov
       uses: codecov/codecov-action@v1

src/overcooked_ai_py/agents/benchmarking.py

@@ -2,12 +2,14 @@
 import numpy as np
 from IPython.display import display
 
-from overcooked_ai_py.utils import save_pickle, load_pickle, cumulative_rewards_from_rew_list, save_as_json, load_from_json, mean_and_std_err, append_dictionaries, merge_dictionaries, rm_idx_from_dict, take_indexes_from_dict
+from overcooked_ai_py.utils import save_pickle, load_pickle, cumulative_rewards_from_rew_list, save_as_json, load_from_json, mean_and_std_err, append_dictionaries, merge_dictionaries, rm_idx_from_dict, take_indexes_from_dict, NumpyArrayEncoder
 from overcooked_ai_py.planning.planners import NO_COUNTERS_PARAMS, MediumLevelPlanner
 from overcooked_ai_py.mdp.layout_generator import LayoutGenerator
 from overcooked_ai_py.agents.agent import AgentPair, CoupledPlanningAgent, RandomAgent, GreedyHumanModel
 from overcooked_ai_py.mdp.overcooked_mdp import OvercookedGridworld, Action, OvercookedState
 from overcooked_ai_py.mdp.overcooked_env import OvercookedEnv
+from overcooked_ai_py.visualization.extract_events import extract_events
+from overcooked_ai_py.visualization.visualization_utils import run_html_in_ipython, run_html_in_web, create_chart_html
 
 
 class AgentEvaluator(object):
@@ -162,7 +164,7 @@ def _check_trajectories_dynamics(trajectories):
 
                 next_state, reward, done, info = simulation_env.step(actions[i])
 
-                assert states[i + 1] == next_state, "States differed (expected vs actual): {}\n\nexpected dict: \t{}\nactual dict: \t{}".format(
+                assert states[i + 1].ids_independent_equal(next_state), "States differed (expected vs actual): {}\n\nexpected dict: \t{}\nactual dict: \t{}".format(
                     simulation_env.display_states(states[i + 1], next_state), states[i+1].to_dict(), next_state.to_dict()
                 )
                 assert rewards[i] == reward, "{} \t {}".format(rewards[i], reward)
@@ -259,19 +261,24 @@ def make_trajectories_json_serializable(trajectories):
         dict_traj['ep_dones'] = [list(lst) for lst in dict_traj['ep_dones']]
         dict_traj['ep_returns'] = [int(val) for val in dict_traj['ep_returns']]
         dict_traj['ep_lengths'] = [int(val) for val in dict_traj['ep_lengths']]
-
-        # NOTE: Currently saving to JSON does not support ep_infos (due to nested np.arrays) or metadata
-        del dict_traj['ep_infos']
+        # NOTE: Currently saving to JSON does not support metadata
+        dict_traj['ep_infos'] = json.loads(json.dumps(dict_traj['ep_infos'], cls=NumpyArrayEncoder))
         del dict_traj['metadatas']
         return dict_traj
 
     @staticmethod
     def load_traj_from_json(filename):
         traj_dict = load_from_json(filename)
+        return AgentEvaluator.load_traj_from_json_obj(traj_dict)
+
+    @staticmethod
+    def load_traj_from_json_obj(obj):
+        # currently ep_infos is not changed back to numpy arrays/tuples from lists
+        traj_dict = copy.deepcopy(obj)
         traj_dict["ep_states"] = [[OvercookedState.from_dict(ob) for ob in curr_ep_obs] for curr_ep_obs in traj_dict["ep_states"]]
         traj_dict["ep_actions"] = [[tuple(tuple(a) if type(a) is list else a for a in j_a) for j_a in ep_acts] for ep_acts in traj_dict["ep_actions"]]
         return traj_dict
-
+    
     ############################
     # TRAJ MANINPULATION UTILS #
     ############################
@@ -379,8 +386,19 @@ def update(t = 1.0):
         display(out, t)
 
     # EVENTS VISUALIZATION METHODS #
-
+        
     @staticmethod
-    def events_visualization(trajs, traj_index):
-        # TODO
-        pass
+    def events_visualization(trajs, traj_index=0, ipython=False, chart_settings=None):
+        """
+        Displays chart with visualization of events (when items pickups happened etc.)
+        ipython - chooses between opening chart
+        in default browser or below ipython cell (for ipython=True)
+        chart_settings - json with various chart settings that overwrittes default ones
+        for more info see create_chart_html function with comments aboutr default chart settings
+        """
+        events = extract_events(trajs, traj_index)
+        html = create_chart_html(events, chart_settings)
+        if ipython:
+            run_html_in_ipython(html)
+        else:
+            run_html_in_web(html)

src/overcooked_ai_py/mdp/overcooked_env.py

@@ -206,6 +206,7 @@ def reset(self):
             "cumulative_shaped_rewards_by_agent": np.array([0] * self.mdp.num_players)
         }
         self.game_stats = {**events_dict, **rewards_dict}
+        self.events_list = []
 
     def is_done(self):
         """Whether the episode is over."""
@@ -235,6 +236,7 @@ def _prepare_info_dict(self, joint_agent_action_info, mdp_infos):
     def _add_episode_info(self, env_info):
         env_info["episode"] = {
             "ep_game_stats": self.game_stats,
+            "events_list": self.events_list,
             "ep_sparse_r": sum(self.game_stats["cumulative_sparse_rewards_by_agent"]),
             "ep_shaped_r": sum(self.game_stats["cumulative_shaped_rewards_by_agent"]),
             "ep_sparse_r_by_agent": self.game_stats["cumulative_sparse_rewards_by_agent"],
@@ -257,7 +259,10 @@ def _update_game_stats(self, infos):
             for idx, event_by_agent in enumerate(event_occurred_by_idx):
                 if event_by_agent:
                     self.game_stats[event_type][idx].append(self.state.timestep)
-
+
+        for event in infos["events_list"]:
+            event["timestep"] = self.state.timestep
+            self.events_list.append(event)
 
     ####################
     # TRAJECTORY LOGIC #