"""Ravens main training script."""
import os
import pickle
import json
import numpy as np
import hydra
from cliport import agents
from cliport import dataset
from cliport import tasks
from cliport.utils import utils
from cliport.environments.environment import Environment
from torch.utils.data import DataLoader
@hydra.main(config_path='./cfg', config_name='eval', version_base="1.2")
def main(vcfg):
# Load train cfg
tcfg = utils.load_hydra_config(vcfg['train_config'])
# Initialize environment and task.
env = Environment(
vcfg['assets_root'],
disp=vcfg['disp'],
shared_memory=vcfg['shared_memory'],
hz=480,
record_cfg=vcfg['record']
)
# Choose eval mode and task.
mode = vcfg['mode']
eval_task = vcfg['eval_task']
print("eval_task!!!", eval_task)
if mode not in {'train', 'val', 'test'}:
raise Exception("Invalid mode. Valid options: train, val, test")
# Load eval dataset.
dataset_type = vcfg['type']
if 'multi' in dataset_type:
ds = dataset.RavensMultiTaskDataset(vcfg['data_dir'],
tcfg,
group=eval_task,
mode=mode,
n_demos=vcfg['n_demos'],
augment=False)
else:
ds = dataset.RavensDataset(os.path.join(vcfg['data_dir'], f"{eval_task}-{mode}"),
tcfg,
n_demos=vcfg['n_demos'],
augment=False)
all_results = {}
name = '{}-{}-n{}'.format(eval_task, vcfg['agent'], vcfg['n_demos'])
# Save path for results.
json_name = f"multi-results-{mode}.json" if 'multi' in vcfg['model_path'] else f"results-{mode}.json"
save_path = vcfg['save_path']
print(f"Save path for results: {save_path}")
if not os.path.exists(save_path):
os.makedirs(save_path)
save_json = os.path.join(save_path, f'{name}-{json_name}')
# Load existing results.
existing_results = {}
if os.path.exists(save_json):
with open(save_json, 'r') as f:
existing_results = json.load(f)
# Make a list of checkpoints to eval.
ckpts_to_eval = list_ckpts_to_eval(vcfg, existing_results)
data_loader = DataLoader(ds, shuffle=False,
pin_memory=False,
num_workers=1 )
# Evaluation loop
print(f"Evaluating: {str(ckpts_to_eval)}")
for ckpt in ckpts_to_eval:
model_file = os.path.join(vcfg['model_path'], ckpt)
if not os.path.exists(model_file) or not os.path.isfile(model_file):
print(f"Checkpoint not found: {model_file}")
continue
elif not vcfg['update_results'] and ckpt in existing_results:
print(f"Skipping because of existing results for {model_file}.")
continue
results = []
mean_reward = 0.0
# Run testing for each training run.
for train_run in range(vcfg['n_repeats']):
# Initialize agent.
utils.set_seed(train_run, torch=True)
agent = agents.names[vcfg['agent']](name, tcfg, data_loader, data_loader)
# Load checkpoint
agent.load(model_file)
print(f"Loaded: {model_file}")
record = vcfg['record']['save_video']
n_demos = vcfg['n_demos']
# Run testing and save total rewards with last transition info.
for i in range(0, n_demos):
print(f'Test: {i + 1}/{n_demos}')
try:
episode, seed = ds.load(i)
except:
print(f"skip bad example {i}")
continue
goal = episode[-1]
total_reward = 0
np.random.seed(seed)
# set task
if 'multi' in dataset_type:
task_name = ds.get_curr_task()
task = tasks.names[task_name]()
print(f'Evaluating on {task_name}')
else:
task_name = vcfg['eval_task']
task = tasks.names[task_name]()
task.mode = mode
env.seed(seed)
env.set_task(task)
obs = env.reset()
info = env.info
reward = 0
# Start recording video (NOTE: super slow)
if record:
video_name = f'{task_name}-{i+1:06d}'
if 'multi' in vcfg['model_task']:
video_name = f"{vcfg['model_task']}-{video_name}"
env.start_rec(video_name)
for _ in range(task.max_steps):
act = agent.act(obs, info, goal)
lang_goal = info['lang_goal']
# print(f'Lang Goal: {lang_goal}')
obs, reward, done, info = env.step(act)
total_reward += reward
# print(f'Total Reward: {total_reward:.3f} | Done: {done}\n')
if done:
break
results.append((total_reward, info))
mean_reward = np.mean([r for r, i in results])
print(f'Mean: {mean_reward} | Task: {task_name} | Ckpt: {ckpt}')
# End recording video
if record:
env.end_rec()
all_results[ckpt] = {
'episodes': results,
'mean_reward': mean_reward,
}
# Save results in a json file.
if vcfg['save_results']:
print("save results to:", save_json)
# Load existing results
if os.path.exists(save_json):
with open(save_json, 'r') as f:
existing_results = json.load(f)
existing_results.update(all_results)
all_results = existing_results
with open(save_json, 'w') as f:
json.dump(all_results, f, indent=4)
def list_ckpts_to_eval(vcfg, existing_results):
ckpts_to_eval = []
# Just the last.ckpt
if vcfg['checkpoint_type'] == 'last':
last_ckpt = 'last.ckpt'
ckpts_to_eval.append(last_ckpt)
# Validation checkpoints that haven't been already evaluated.
elif vcfg['checkpoint_type'] == 'val_missing':
checkpoints = sorted([c for c in os.listdir(vcfg['model_path']) if "steps=" in c])
ckpts_to_eval = [c for c in checkpoints if c not in existing_results]
# Find the best checkpoint from validation and run eval on the test set.
elif vcfg['checkpoint_type'] == 'test_best':
result_jsons = [c for c in os.listdir(vcfg['results_path']) if "results-val" in c]
if 'multi' in vcfg['model_task']:
result_jsons = [r for r in result_jsons if "multi" in r]
else:
result_jsons = [r for r in result_jsons if "multi" not in r]
if len(result_jsons) > 0:
result_json = result_jsons[0]
with open(os.path.join(vcfg['results_path'], result_json), 'r') as f:
eval_res = json.load(f)
best_checkpoint = 'last.ckpt'
best_success = -1.0
for ckpt, res in eval_res.items():
if res['mean_reward'] > best_success:
best_checkpoint = ckpt
best_success = res['mean_reward']
print(best_checkpoint)
ckpt = best_checkpoint
ckpts_to_eval.append(ckpt)
else:
print("No best val ckpt found. Using last.ckpt")
ckpt = 'last.ckpt'
ckpts_to_eval.append(ckpt)
# Load a specific checkpoint with a substring e.g: 'steps=10000'
else:
print(f"Looking for: {vcfg['checkpoint_type']}")
checkpoints = [c for c in os.listdir(vcfg['model_path']) if vcfg['checkpoint_type'] in c]
checkpoint = checkpoints[0] if len(checkpoints) > 0 else ""
ckpt = checkpoint
ckpts_to_eval.append(ckpt)
print("ckpts_to_eval:", ckpts_to_eval)
return ckpts_to_eval
if __name__ == '__main__':
main()