RonyForAI/Mirage_DB_RL

1import torch
2from Mirage_RL.training.agent import Agent
3from Mirage_RL.client import QueryClient
4from Mirage_RL.models import QueryAction
5
6agent = Agent(num_tables=3)
7episodes = 100
8
9with QueryClient(base_url="http://localhost:8000").sync() as env:
10    for episode in range(episodes):
11        result = env.reset()
12        obs = result.observation
13
14        state = agent.encode_state(obs)
15        total_reward = 0
16        done = False
17
18        while not done:
19            # select action
20            (table, join, index), action_id = agent.select_action(obs)
21
22            action = QueryAction(
23                next_table=table,
24                join_type=join,
25                use_index=index
26            )
27
28            result = env.step(action)
29
30            next_obs = result.observation
31            reward = result.reward
32            done = result.done
33
34            next_state = agent.encode_state(next_obs)
35
36            # train
37            agent.train_step(state, action_id, reward, next_state, done)
38
39            state = next_state
40            obs = next_obs
41            total_reward += reward
42
43        print(f"Episode {episode:>3} | Total Reward: {total_reward:>8.2f} | Epsilon: {agent.epsilon:.3f}")