-
Notifications
You must be signed in to change notification settings - Fork 0
/
eval.py
40 lines (33 loc) · 1.52 KB
/
eval.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
from robocasa.environments import ALL_KITCHEN_ENVIRONMENTS
from robocasa.utils.dataset_registry import SINGLE_STAGE_TASK_DATASETS, MULTI_STAGE_TASK_DATASETS
from robocasa.utils.dataset_registry import get_ds_path
from robocasa.utils.env_utils import create_env, run_random_rollouts
import numpy as np
"""
Select a random task (that comes with an accompanying dataset) to run rollouts for.
Alternatively, sample *any* kitchen task in RoboCasa by replacing the following line with
env_name = np.random.choice(list(ALL_KITCHEN_ENVIRONMENTS))
"""
# env_name = np.random.choice(
# list(SINGLE_STAGE_TASK_DATASETS)
# )
print(list(MULTI_STAGE_TASK_DATASETS))
env_name = "PnPCounterToStove"
"""
['PnPCounterToCab', 'PnPCabToCounter', 'PnPCounterToSink', 'PnPSinkToCounter', 'PnPCounterToMicrowave',
'PnPMicrowaveToCounter', 'PnPCounterToStove', 'PnPStoveToCounter', 'OpenSingleDoor', 'CloseSingleDoor',
'OpenDoubleDoor', 'CloseDoubleDoor', 'OpenDrawer', 'CloseDrawer', 'TurnOnSinkFaucet', 'TurnOffSinkFaucet',
'TurnSinkSpout', 'TurnOnStove', 'TurnOffStove', 'CoffeeSetupMug', 'CoffeeServeMug', 'CoffeePressButton',
'TurnOnMicrowave', 'TurnOffMicrowave', 'NavigateKitchen']
"""
"""
['ArrangeVegetables', 'MicrowaveThawing', 'RestockPantry', 'PreSoakPan', 'PrepareCoffee']
"""
print(env_name)
# seed environment as needed. set seed=None to run unseeded
env = create_env(env_name=env_name, seed=0)
# run rollouts with random actions and save video
info = run_random_rollouts(
env, num_rollouts=1, num_steps=800, video_path="/tmp/test.mp4"
)
print(info)