Fix reproduce episode

[vkakerbeck]

The make_detailed_follow_up_configs.py script was not working anymore after this PR #153 The changes here should fix this.

Loading the config is a bit more involved now. Previously, the script would print instructions for running the follow-up config, which was simply python run.py -e {follow_up_name}. Now, as far as I understand, one needs to add the follow_up_name to one of the classes in names.py and then add

from benchmarks.configs.follow_ups import CONFIGS as FOLLOW_UP_CONFIGS
experiments = YcbExperiments(
    follow_up_name=FOLLOW_UP_CONFIGS["follow_up_name"],
)

to the corresponding experiment config script (here ycb_experiments.py). This is a lot more complicated.
@tristanls is there a simpler way that I am missing? Otherwise, I will update the instructions in make_detailed_follow_up_configs.py.

I also noticed that the predefined action sequence wasn't applied properly, and a wandb handler was used that doesn't have all its dependencies in the environment setup, so I fixed those two.

I am looking into why I am still not getting the exact episode reproduced, hence the draft PR.

[tristanls]

issue: I think setup_env should go into benchmarks/make_detailed_follow_up_configs.py.

It looks like benchmarks/make_detailed_follow_up_configs.py is analogous to benchmarks/run.py and benchmarks/run_parallel.py. I added a comment to benchmarks/make_detailed_follow_up_configs.py as to where I think setup_env should go.

[tristanls]

issue: I think calling setup_env() should happen before this line. In run.py and run_parallel.py we call setup_env() before __main__.

[tristanls]

Correction, the setup_env() should go before the follow_up_configs import like so:

from tbp.monty.frameworks.run_env import setup_env

setup_env()

from tbp.monty.frameworks.utils.follow_up_configs import (  # noqa: E402
    create_eval_config_multiple_episodes,
    recover_output_dir,
)

[tristanls]

The above change gives me:

(tbp.monty) me@revisionist tbp.monty % python benchmarks/make_detailed_follow_up_configs.py -h
MONTY_LOGS not set. Using default directory: ~/tbp/results/monty/
MONTY_MODELS not set. Using default directory: ~/tbp/results/monty/pretrained_models/
MONTY_DATA not set. Using default directory: /Users/me/tbp/data/
WANDB_DIR not set. Using default directory: ~/tbp/results/monty/
usage: make_detailed_follow_up_configs.py [-h]
                                          [-e {world_image_from_stream_on_scanned_model,world_image_on_scanned_model,dark_world_image_on_scanned_model,bright_world_image_on_scanned_model,hand_intrusion_world_image_on_scanned_model,multi_object_world_image_on_scanned_model,randrot_noise_sim_on_scan_monty_world,supervised_pre_training_base,supervised_pre_training_5lms,supervised_pre_training_5lms_all_objects,only_surf_agent_training_10obj,only_surf_agent_training_10simobj,only_surf_agent_training_allobj,only_surf_agent_training_numenta_lab_obj,base_config_10distinctobj_dist_agent,base_config_10distinctobj_surf_agent,randrot_noise_10distinctobj_dist_agent,randrot_noise_10distinctobj_dist_on_distm,randrot_noise_10distinctobj_surf_agent,randrot_10distinctobj_surf_agent,randrot_noise_10distinctobj_5lms_dist_agent,base_10simobj_surf_agent,randrot_noise_10simobj_surf_agent,randrot_noise_10simobj_dist_agent,randomrot_rawnoise_10distinctobj_surf_agent,base_10multi_distinctobj_dist_agent,surf_agent_unsupervised_10distinctobj,surf_agent_unsupervised_10distinctobj_noise,surf_agent_unsupervised_10simobj,base_77obj_dist_agent,base_77obj_surf_agent,randrot_noise_77obj_surf_agent,randrot_noise_77obj_dist_agent,randrot_noise_77obj_5lms_dist_agent}]
                                          [-i EPISODES [EPISODES ...]] [-n NAME]

optional arguments:
  -h, --help            show this help message and exit
  -e {world_image_from_stream_on_scanned_model,world_image_on_scanned_model,dark_world_image_on_scanned_model,bright_world_image_on_scanned_model,hand_intrusion_world_image_on_scanned_model,multi_object_world_image_on_scanned_model,randrot_noise_sim_on_scan_monty_world,supervised_pre_training_base,supervised_pre_training_5lms,supervised_pre_training_5lms_all_objects,only_surf_agent_training_10obj,only_surf_agent_training_10simobj,only_surf_agent_training_allobj,only_surf_agent_training_numenta_lab_obj,base_config_10distinctobj_dist_agent,base_config_10distinctobj_surf_agent,randrot_noise_10distinctobj_dist_agent,randrot_noise_10distinctobj_dist_on_distm,randrot_noise_10distinctobj_surf_agent,randrot_10distinctobj_surf_agent,randrot_noise_10distinctobj_5lms_dist_agent,base_10simobj_surf_agent,randrot_noise_10simobj_surf_agent,randrot_noise_10simobj_dist_agent,randomrot_rawnoise_10distinctobj_surf_agent,base_10multi_distinctobj_dist_agent,surf_agent_unsupervised_10distinctobj,surf_agent_unsupervised_10distinctobj_noise,surf_agent_unsupervised_10simobj,base_77obj_dist_agent,base_77obj_surf_agent,randrot_noise_77obj_surf_agent,randrot_noise_77obj_dist_agent,randrot_noise_77obj_5lms_dist_agent}, -exp {world_image_from_stream_on_scanned_model,world_image_on_scanned_model,dark_world_image_on_scanned_model,bright_world_image_on_scanned_model,hand_intrusion_world_image_on_scanned_model,multi_object_world_image_on_scanned_model,randrot_noise_sim_on_scan_monty_world,supervised_pre_training_base,supervised_pre_training_5lms,supervised_pre_training_5lms_all_objects,only_surf_agent_training_10obj,only_surf_agent_training_10simobj,only_surf_agent_training_allobj,only_surf_agent_training_numenta_lab_obj,base_config_10distinctobj_dist_agent,base_config_10distinctobj_surf_agent,randrot_noise_10distinctobj_dist_agent,randrot_noise_10distinctobj_dist_on_distm,randrot_noise_10distinctobj_surf_agent,randrot_10distinctobj_surf_agent,randrot_noise_10distinctobj_5lms_dist_agent,base_10simobj_surf_agent,randrot_noise_10simobj_surf_agent,randrot_noise_10simobj_dist_agent,randomrot_rawnoise_10distinctobj_surf_agent,base_10multi_distinctobj_dist_agent,surf_agent_unsupervised_10distinctobj,surf_agent_unsupervised_10distinctobj_noise,surf_agent_unsupervised_10simobj,base_77obj_dist_agent,base_77obj_surf_agent,randrot_noise_77obj_surf_agent,randrot_noise_77obj_dist_agent,randrot_noise_77obj_5lms_dist_agent}, --experiment {world_image_from_stream_on_scanned_model,world_image_on_scanned_model,dark_world_image_on_scanned_model,bright_world_image_on_scanned_model,hand_intrusion_world_image_on_scanned_model,multi_object_world_image_on_scanned_model,randrot_noise_sim_on_scan_monty_world,supervised_pre_training_base,supervised_pre_training_5lms,supervised_pre_training_5lms_all_objects,only_surf_agent_training_10obj,only_surf_agent_training_10simobj,only_surf_agent_training_allobj,only_surf_agent_training_numenta_lab_obj,base_config_10distinctobj_dist_agent,base_config_10distinctobj_surf_agent,randrot_noise_10distinctobj_dist_agent,randrot_noise_10distinctobj_dist_on_distm,randrot_noise_10distinctobj_surf_agent,randrot_10distinctobj_surf_agent,randrot_noise_10distinctobj_5lms_dist_agent,base_10simobj_surf_agent,randrot_noise_10simobj_surf_agent,randrot_noise_10simobj_dist_agent,randomrot_rawnoise_10distinctobj_surf_agent,base_10multi_distinctobj_dist_agent,surf_agent_unsupervised_10distinctobj,surf_agent_unsupervised_10distinctobj_noise,surf_agent_unsupervised_10simobj,base_77obj_dist_agent,base_77obj_surf_agent,randrot_noise_77obj_surf_agent,randrot_noise_77obj_dist_agent,randrot_noise_77obj_5lms_dist_agent}
                        Name of the experiment you want to repeat episodes from (default: None)
  -i EPISODES [EPISODES ...], --idx EPISODES [EPISODES ...], --episodes EPISODES [EPISODES ...]
                        List of int, indices of episodes you want to re-run. Index lookup will be based on file names from reproducibility logger, e.g. reproduce_episode_data/eval_episode_2_actions.jsonl (default: [])
  -n NAME, --name NAME  name of follow up config that will be appended to parent name (default: rerun)

[vkakerbeck]

Agreed. That's where I had put it before this commit but had then moved it to pass ruff import order requirements. I moved it back now with the # noqa: E402 instead, as you suggested.

[tristanls]

@vkakerbeck, I think we can generate the code so that those instructions remain the same.

I think we can create and then append to benchmarks/configs/follow_ups/names.py by writing a new line for each follow_up_name, e.g., if I created my_follow_up1, we'd append NAMES.append("my_follow_up1") to benchmarks/configs/follow_ups/names.py:

# benchmarks/configs/follow_ups/names.py
NAMES = []

NAMES.append("my_follow_up1")

Then, if I created another_follow_up, we'd append another line to that file:

# benchmarks/config/follow_ups/names.py
NAMES = []

NAMES.append("my_follow_up1")
NAMES.append("another_follow_up")

This benchmarks/configs/follow_ups/names.py is less involved because we are not adding the safety mechanism to remind people to change code in two places. This code is autogenerated, so we don't need that safety. It also makes it easier to implement.

The configs would continue to be written as they currently are.

Then, we'd update benchmarks/configs/load.py as highlighted below:

# benchmarks/configs/load.py
...

from __future__ import annotations

from dataclasses import fields

from benchmarks.configs.names import (
    MontyWorldExperiments,
    MontyWorldHabitatExperiments,
    MyExperiments,
    PretrainingExperiments,
    YcbExperiments,
)
####### Add this import (would also add this to benchmarks/run.py and benchmarks/run_parallel.py
from benchmarks.configs.follow_ups.names import NAMES as follow_up_experiment_names
#######

__all__ = ["load_configs", "load_config"]

...

def select_config(experiment: str) -> dict:
    """Select and import the configuration group for the given experiment.

    Args:
        experiment: The name of the experiment to import the configuration for.

    Returns:
        The imported configuration group for the given experiment.

    Note:
        We import the configurations selectively to avoid importing uninstalled
        dependencies. For example, if someone does not install a specific simulator,
        then do not import the corresponding configuration group.
    """
    monty_world_experiment_names = [
        field.name for field in fields(MontyWorldExperiments)
    ]
    monty_world_habitat_experiment_names = [
        field.name for field in fields(MontyWorldHabitatExperiments)
    ]
    pretraining_experiment_names = [
        field.name for field in fields(PretrainingExperiments)
    ]
    ycb_experiment_names = [field.name for field in fields(YcbExperiments)]
    my_experiment_names = [field.name for field in fields(MyExperiments)]

    if experiment in monty_world_experiment_names:
        from benchmarks.configs.monty_world_experiments import (
            CONFIGS as MONTY_WORLD,
        )

        return MONTY_WORLD
    elif experiment in monty_world_habitat_experiment_names:
        from benchmarks.configs.monty_world_habitat_experiments import (
            CONFIGS as MONTY_WORLD_HABITAT,
        )

        return MONTY_WORLD_HABITAT
    elif experiment in pretraining_experiment_names:
        from benchmarks.configs.pretraining_experiments import (
            CONFIGS as PRETRAININGS,
        )

        return PRETRAININGS
    elif experiment in ycb_experiment_names:
        from benchmarks.configs.ycb_experiments import CONFIGS as YCB

        return YCB
    elif experiment in my_experiment_names:
        from benchmarks.configs.my_experiments import CONFIGS as MY_EXPERIMENTS

        return MY_EXPERIMENTS
###### Select the config as usual
    elif experiment in follow_up_experiment_names:
        from benchmarks.configs.follow_ups import CONFIGS as FOLLOW_UP_EXPERIMENTS

        return FOLLOW_UP_EXPERIMENTS
######