<jupyter_start><jupyter_text>Exploring TEACh Data<jupyter_code>import os
import sys
import json
import copy

sys.path.append("../../")
from teach.dataset.definitions import Definitions
from teach.dataset.dataset import Dataset
from teach.dataset.actions import Action_Keyboard, Action_ObjectInteraction
# Edit data directory if changed when using `teach_download`
data_dir = "/tmp/teach-dataset"<jupyter_output><empty_output><jupyter_text>Definitions Instantiate a `Definitions` object to access various definitions, mappings of agent IDs and actions to names, as well as task definitions. The code uses `Driver` when referring to the `Follower` in the paper.<jupyter_code>definitions = Definitions(version="2.0")
print("Agent IDs to agents: ", definitions.map_agents_id2info)
print("Status IDs to names: ", definitions.map_status_id2name)<jupyter_output>Agent IDs to agents:  OrderedDict([(0, OrderedDict([('agent_name', 'Commander'), ('agent_type', 0)])), (1, OrderedDict([('agent_name', 'Driver'), ('agent_type', 1)]))])
Status IDs to names:  OrderedDict([(0, 'Success'), (1, 'Failure')])<jupyter_text>Display mappings of action IDs to action names. Note that only a subset of these are used in TEACh data. Note that `definitions.map_tasks_name2info` ends up being more useful when trying to access actions by name.<jupyter_code>print("Action IDs to names:")
for action_id, action in definitions.map_actions_id2info.items():
    print("\t ", action_id, ":", action["action_name"])<jupyter_output>Action IDs to names:
	  0 : Stop
	  1 : Move to
	  2 : Forward
	  3 : Backward
	  4 : Turn Left
	  5 : Turn Right
	  6 : Look Up
	  7 : Look Down
	  8 : Pan Left
	  9 : Pan Right
	  10 : Move Up
	  11 : Move Down
	  12 : Double Forward
	  13 : Double Backward
	  300 : Navigation
	  200 : Pickup
	  201 : Place
	  202 : Open
	  203 : Close
	  204 : ToggleOn
	  205 : ToggleOff
	  206 : Slice
	  207 : Dirty
	  208 : Clean
	  209 : Fill
	  210 : Empty
	  211 : Pour
	  212 : Break
	  400 : BehindAboveOn
	  401 : BehindAboveOff
	  500 : OpenProgressCheck
	  501 : SelectOid
	  502 : SearchObject
	  100 : Text
	  101 : Speech
	  102 : Beep<jupyter_text>Tasks are also most convenient to access by name via `definitions.map_tasks_name2info` but can be accessed via ID using `definitions.map_tasks_id2info`. The values of both of these dictionaries are of type `Task_THOR`.  When a `Definitions` object is instantiated, all tasks defined under `src/teach/meta_data_files/task_definitions` get loaded. The Task Definition Language is explained in Appendix F of the [TEACh paper](https://arxiv.org/pdf/2110.00534.pdf). To create a new task, create a new JSON file under `src/teach/meta_data_files/task_definitions`. Each task needs to have a unique `task_id` and `task_name`. Tasks can be referenced in other tasks by their `task_name`. After creating a new task, test that it can be loaded any any inter-task dependencies can be resolved by instantiating a `Definitions` object.The following code snippet demonstrates how to print a few task details. Note that `n` (where `n` is a number) indicates a variable.<jupyter_code>print("Task details by name:")
print("Task name".ljust(33, " "), "Task ID".ljust(10, " "), "Num task params".ljust(20, " "), "Task component names")
for task_name, task in definitions.map_tasks_name2info.items():
    print(
        task_name.ljust(35, " "),
        str(task.task_id).ljust(15, " "),
        str(task.task_nparams).ljust(10, " "),
        str(list(task.components.keys())),
    )<jupyter_output>Task details by name:
Task name                         Task ID    Num task params      Task component names
Breakfast                           301             14         ['coffee', 'toast', 'potatoes', 'apple', 'sandwich', 'salad', 'serving_spot']
Boil X                              112             1          ['boiled_#0']
Basic Kitchen Tasks                 401             0          ['coffee', 'toast', 'omelette', 'spatula', 'drawer']
Tutorial                            201             0          ['coffee', 'potato']
Water Plant                         114             0          ['water_plant']
Sliced X                            104             1          ['#0', 'knife']
Basic Bathroom Tasks                403             0          ['faucets', 'candles', 'soap', 'counter']
Sandwich                            302             2          ['toast', 'lettuce', 'tomato', 'plate']
Put All X On Y                      110             3          ['#0', '#2']
Poach Egg                      [...]<jupyter_text>Gameplay SessionsGameplay sessions are stored in `json` files. The `games` subdirectory consists of one subdirectory per split each containing game files of that split. When loaded, these are dictionaries and for many purposes, it is sufficient to analyze the dictionaries. Some examples:<jupyter_code>f = os.path.join(data_dir, "games/train/7d2a79f43e605c36_1657.game.json")
with open(f) as h:
    game_dict = json.load(h)
print(game_dict.keys())<jupyter_output>dict_keys(['version', 'task_type', 'comments', 'definitions', 'tasks'])<jupyter_text>While the game dictionary contains other keys, the important one is `tasks`. `version`, `task_type` and `comments` are dataset-specific metadata, and `definitions` contains the version of the `Definitions` object used to collect the data. However, all games in the subdirectory `games` have been verified to be replayable and resulting in task success using the current (released) version of the `Definitions` object. `tasks` is always a list of length 1 in this dataset.<jupyter_code>print(game_dict["tasks"][0].keys())<jupyter_output>dict_keys(['task_id', 'task_name', 'task_params', 'task_nparams', 'task_anchor_object', 'desc', 'components', 'relations', 'comments', 'episodes'])<jupyter_text>This is a dictionary that can be converted to a `Task_THOR` object. All keys except `episodes` are associated with the task definition and can be better understood by reading Appendix F of the [TEACh paper](https://arxiv.org/pdf/2110.00534.pdf). For all game files in this dataset `game_dict['tasks'][0]['episodes']` will be a list of length 1 and `game_dict['tasks'][0]['episodes'][0]` contains the actual sequence of actions taken in the episode.<jupyter_code>print(game_dict["tasks"][0]["episodes"][0].keys())<jupyter_output>dict_keys(['episode_id', 'world', 'world_type', 'commander_embodied', 'initial_state', 'interactions', 'final_state'])<jupyter_text>Episodes are used to store the initial and final simulator state, as well as the sequence of actions taken in a gameplay session. The components of an episode are:* `episode_id` - A unique id* `world_type` - Type of room which is one of `Kitchen`, `Bedroom`, `Bathroom` and `Living room` * `world` - ID of the specific AI2-THOR floor plan used for this gameplay session* `commander_embodied` - False for all TEACh games* `initial_state`, `final_state` - Dictionaries consisting of the initial and final state of the world including    * `time_start` -     * `agents` - Position and orientation of each agent/ camera at start and end of episode    * `objects` - A list of the state of all objects at the start and end of the episode. Each object is represented by a dictionary whose keys are property names and values are property values.    * `custom_object_metadata` - A dictionary to track custom properties in our codebase that are not present in AI2-THOR. This is a dictionary with AI2-THOR objectId as key and a dictionary of (custom_property_name, custom_property_value) pairs as values* `interactions` - An ordered list of interactions that occurred in the environment, each represented by a dictionary of    * `agent_id` - The agent that took the action    * `action_id` - Which action was taken    * `time_start` - Duration of time between start of episode and when this action started    * `duration` - Duration of time (in sec) taken to execute this action    * `success` - 1 if the action was successfully executed during data collection and 0 otherwise. An example of a case where `success` might be 0 is if the human annotator tried to pick up an object from too far away     * Action specific keys. Some examples include        * `utterance` for a `Text` action - Stores the text value of the utterance made        * `pose_delta` and `pose` for a navigation action        Code snippet to print out the sequence of actions taken in an episode:<jupyter_code>def print_actions_from_game_dict(game_dict, definitions):
    interactions = game_dict["tasks"][0]["episodes"][0]["interactions"]
    print(
        "Time Start",
        "Action Success".ljust(15, " "),
        "Agent".ljust(15, " "),
        "Action".ljust(20, " "),
        "Utterance text / Object ID / Object X, Y",
    )
    for interaction in interactions:
        output_str = "".rjust(2, " ")
        output_str += ("%.2f" % interaction["time_start"]).ljust(15, " ")
        output_str += str(interaction["success"]).ljust(10, " ")
        output_str += definitions.map_agents_id2info[interaction["agent_id"]]["agent_name"].ljust(15, " ")
        output_str += definitions.map_actions_id2info[interaction["action_id"]]["action_name"].ljust(20, " ")
        if "utterance" in interaction:
            output_str += interaction["utterance"]
        elif "oid" in interaction and interaction["oid"] is not None:
            output_str += interaction["oid"]
        elif "x" in interaction and "y" in interaction:
            output_str += "(" + str(interaction["x"]) + ", " + str(interaction["y"]) + ")"
        print(output_str)
print_actions_from_game_dict(game_dict, definitions)<jupyter_output>Time Start Action Success  Agent           Action               Utterance text / Object ID / Object X, Y
  15.29          0         Commander      OpenProgressCheck   
  27.85          1         Commander      Text                I need the newspaper to be placed on a single table.
  29.49          1         Commander      SelectOid           
  39.11          1         Driver         Text                what should i do
  61.21          1         Driver         Pan Left            
  61.59          1         Driver         Pan Left            
  61.84          1         Driver         Pan Left            
  62.12          1         Commander      Text                I need the newspaper placed on a single table.
  70.16          1         Driver         Pickup              Newspaper|-04.15|+00.36|-02.48
  87.74          1         Driver         Place               CoffeeTable|-02.47|+00.00|-02.49
  92.55          1         Commander      OpenProgressCheck<jupyter_text>Note that for all object interactions, the relative coordinates of the object on the agent's egocentric image are available in `interaction['x'], interaction['y']`. In the cases where the wrapper was able to resolve these to an object ID using the segmentation frame, we also have the ID of the object interacted with in `interaction['oid']` but if the wrapper was forced to backoff to raycasting, then this is not available.   It is also possible to import a game file into a `Dataset` object as follows.<jupyter_code>f = os.path.join(data_dir, "games/train/7d2a79f43e605c36_1657.game.json")
game = Dataset.import_json(f)<jupyter_output><empty_output><jupyter_text>The following is how the code snippet to print out the same action info would look using the object oriented representation:<jupyter_code>def print_actions_from_game_as_dataset(game, definitions):
    interactions = game.tasks[0].episodes[0].interactions
    print(
        "Time Start",
        "Action Success".ljust(15, " "),
        "Agent".ljust(15, " "),
        "Action".ljust(20, " "),
        "Utterance text / Object ID / Object X, Y",
    )
    for interaction in interactions:
        output_str = "".rjust(2, " ")
        output_str += ("%.2f" % interaction.time_start).ljust(15, " ")
        output_str += str(interaction.status).ljust(10, " ")
        output_str += definitions.map_agents_id2info[interaction.agent_id]["agent_name"].ljust(15, " ")
        output_str += definitions.map_actions_id2info[interaction.action.action_id]["action_name"].ljust(20, " ")
        if isinstance(interaction.action, Action_Keyboard):
            output_str += interaction.action.utterance
        if isinstance(interaction.action, Action_ObjectInteraction):
            if interaction.action.oid is None:
                output_str += "(" + str(interaction.action.x) + ", " + str(interaction.action.y) + ")"
            else:
                output_str += interaction.action.oid
        print(output_str)
print_actions_from_game_as_dataset(game, definitions)<jupyter_output>Time Start Action Success  Agent           Action               Utterance text / Object ID / Object X, Y
  15.29          None      Commander      OpenProgressCheck   
  27.85          None      Commander      Text                I need the newspaper to be placed on a single table.
  29.49          None      Commander      SelectOid           
  39.11          None      Driver         Text                what should i do
  61.21          None      Driver         Pan Left            
  61.59          None      Driver         Pan Left            
  61.84          None      Driver         Pan Left            
  62.12          None      Commander      Text                I need the newspaper placed on a single table.
  70.16          None      Driver         Pickup              Newspaper|-04.15|+00.36|-02.48
  87.74          None      Driver         Place               CoffeeTable|-02.47|+00.00|-02.49
  92.55          None      Commander      OpenProgressCheck<jupyter_text>Note that while the object oriented representation of the game can be manipulated more easily in the code, the task of the game does not get perfectly loaded. Specifically, when loading a game file, no attempt is made to resolve components of tasks that are themselves tasks. Additionally, the final state does not get loaded. The following code snippet shows how to check whether the task associated with a gameplay session is complete at the final state, by directly loading the game json file as a dictionary.<jupyter_code>definitions = Definitions(version="2.0")
f = os.path.join(data_dir, "games/train/7d2a79f43e605c36_1657.game.json")
with open(f) as h:
    game_dict = json.load(h)
game_task = game_dict["tasks"][0]
task_to_check = copy.deepcopy(
    definitions.map_tasks_name2info[game_task["task_name"]]
)  # Copying is important if you're sharing a definitions object across calls
task_to_check.task_params = game_task["task_params"]
final_state_objects = game_dict["tasks"][0]["episodes"][0]["final_state"]["objects"]
task_check_output = task_to_check.check_episode_progress(final_state_objects)
print(task_check_output["success"])<jupyter_output>True<jupyter_text>EDH InstancesEDH instances are stored in `json` files. The `edh_instances` subdirectory consists of one subdirectory per split each containing EDH instances of that split. EDH instances do not have a corresponding object oriented representation and need to be manipulated as dictionaries.<jupyter_code>f = os.path.join(data_dir, "edh_instances/train/7d2a79f43e605c36_1657.edh0.json")
with open(f) as h:
    edh_instance = json.load(h)
print(edh_instance.keys())<jupyter_output>dict_keys(['future_subgoals', 'pred_start_idx', 'expected_init_goal_conditions_satisfied', 'interactions', 'init_state_diff', 'expected_init_goal_conditions_total', 'expected_init_success', 'instance_id', 'driver_image_history', 'dialog_history', 'state_changes', 'driver_actions_future', 'driver_images_future', 'final_state_diff', 'game_id', 'driver_action_history', 'dialog_history_cleaned', 'history_subgoals'])
{'fraction_non_alphanumeric': 0.0947227703864027, 'fraction_numerical': 0.027882105799223376, 'mean_word_length': 3.098617271067049, 'pattern_counts': {'":': 1, '<': 39, '<?xml version=': 0, '>': 39, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 21, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '12', 'dir': 'jupyter-structured-clean-dedup', 'id': '635365', 'license': 'MIT', 'n_tokens_mistral': 4748, 'n_tokens_neox': 4414, 'n_words': 1540, 'path': 'src/teach/analysis/teach_data_exploration.ipynb', 'repo_name': 'aliang8/teach', 'size': '15709'}