File size: 5,637 Bytes
ff66cf3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182

from hydra.core.global_hydra import GlobalHydra
import gradio as gr
import os
import hydra
import random

import re
import openai
import IPython
import time
import pybullet as p
import traceback
from datetime import datetime
from pprint import pprint
import cv2
import re
import random
import json

from gensim.agent import Agent
from gensim.critic import Critic
from gensim.sim_runner import SimulationRunner
from gensim.memory import Memory
from gensim.utils import set_gpt_model, clear_messages


class DemoRunner:

    def __init__(self):
        self._env = None
        GlobalHydra.instance().clear()
        hydra.initialize(version_base="1.2", config_path='cliport/cfg')
        self._cfg = hydra.compose(config_name="data")

    def setup(self, api_key):
        cfg = self._cfg
        openai.api_key = api_key
        cfg['model_output_dir'] = 'temp'
        cfg['prompt_folder'] = 'bottomup_task_generation_prompt'
        set_gpt_model(cfg['gpt_model'])
        cfg['load_memory'] = True
        cfg['use_template'] = True
        cfg['task_description_candidate_num'] = 10
        cfg['record']['save_video'] = True

        print("cfg = ", cfg)
        memory = Memory(cfg)
        agent = Agent(cfg, memory)
        critic = Critic(cfg, memory)
        self.simulation_runner = SimulationRunner(cfg, agent, critic, memory)

        info = '### Set up '

        return info

    def setup_top_down(self, api_key, target_task_name):
        cfg = self._cfg
        openai.api_key = api_key
        cfg['model_output_dir'] = 'temp'
        cfg['prompt_folder'] = 'topdown_task_generation_prompt'
        set_gpt_model(cfg['gpt_model'])
        cfg['load_memory'] = True
        cfg['use_template'] = True
        cfg['target_task_name'] = target_task_name
        cfg['task_description_candidate_num'] = 10
        cfg['record']['save_video'] = True

        print("cfg = ", cfg)
        memory = Memory(cfg)
        agent = Agent(cfg, memory)
        critic = Critic(cfg, memory)
        self.simulation_runner = SimulationRunner(cfg, agent, critic, memory)

        info = '### Set up '

        return info

    def run(self, instruction, progress):
        cfg = self._cfg
        cfg['target_task_name'] = instruction

        # self._env.cache_video = []
        self.simulation_runner._md_logger = ''
        # progress(0.2)
        yield "Task Generating ==>", None, None
        yield from self.simulation_runner.task_creation()
        yield from self.simulation_runner.simulate_task()

    def run_example(self):
        cfg = self._cfg

        # self._env.cache_video = []
        self.simulation_runner._md_logger = ''
        # progress(0.2)
        yield "Task Generating ==>", None, None
        yield from self.simulation_runner.example_task_creation()
        yield from self.simulation_runner.simulate_task()


def setup(api_key, option_choice, target_task_name):
    print(option_choice)
    if not api_key:
        return 'Please enter your OpenAI API key!', None

    if option_choice is None:
        return 'Please choose the mode!', None
    demo_runner = DemoRunner()

    if option_choice == 'top-down':
        info = demo_runner.setup_top_down(api_key, target_task_name) + option_choice
    elif option_choice == 'bottom-up':
        info = demo_runner.setup(api_key) + option_choice
    else:
        raise NotImplementedError
    return info, demo_runner



def run(instruction, demo_runner, progress=gr.Progress()):
    yield from demo_runner.run(instruction, progress=progress)

def run_example():
    demo_runner = DemoRunner()
    demo_runner.setup(1)
    yield from demo_runner.run_example()


if __name__ == '__main__':
    os.environ['GENSIM_ROOT'] = os.getcwd()
    with open('README.md', 'r') as f:
        for _ in range(12):
            next(f)
        readme_text = f.read()

    with gr.Blocks() as demo:
        state = gr.State(None)

        gr.Markdown(readme_text)
        gr.Markdown('# Interactive Demo')
        with gr.Row():
            with gr.Column():


                btn_example_run = gr.Button("Run Example (OpenAI API Key not required)")
                with gr.Row():
                    inp_api_key = gr.Textbox(label='OpenAI API Key (this is not stored anywhere)', lines=1)

                option_choice = gr.Radio(["bottom-up", "top-down"], label="Which mode?", interactive=True)
                inp_instruction = gr.Textbox(label='Target Task Name (if top-down)', lines=1)
                info_setup = gr.Markdown(label='Setup Info')
                btn_setup = gr.Button("Setup/Reset Simulation")
                btn_run = gr.Button("Run (this may take 30+ seconds)")
            # with gr.Column():

        with gr.Row():
            with gr.Column(scale=1, min_width=600):
                progress = gr.Markdown(label='Progress')
                generated_task = gr.Markdown(label='Generated Task')
                generated_asset = gr.Markdown(label='Generated Asset')
                generated_code = gr.Code(label='Generated Code',  language="python", interactive=True)
                video_run = gr.Video(label='Video of Last Instruction')
        btn_setup.click(
            setup,
            inputs=[inp_api_key, option_choice, inp_instruction],
            outputs=[info_setup, state]
        )
        btn_run.click(
            run,
            inputs=[inp_instruction, state],
            outputs=[progress, generated_code, video_run]
        )

        btn_example_run.click(
            run_example,
            inputs=[],
            outputs=[progress, generated_code, video_run]
        )


    demo.queue().launch(show_error=True)