-

Notifications

You must be signed in to change notification settings

- Fork 177

Expand file tree

/

Copy pathtask.py

More file actions

111 lines (84 loc) · 3.27 KB

/

Copy pathtask.py

File metadata and controls

111 lines (84 loc) · 3.27 KB

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

from abc import ABC, abstractmethod

from typing import Tuple

import numpy as np

import playwright.sync_api

class AbstractBrowserTask(ABC):

"""

Abstract class for browsergym tasks.

"""

@classmethod

def get_task_id(cls):

raise NotImplementedError

def __init__(self, seed: int) -> None:

# initiate a random number generator

self.random = np.random.RandomState(seed)

# task properties, will be used to set up the browsergym environment

# default values, can be overriden in children classes

self.viewport = {"width": 1280, "height": 720}

self.slow_mo = 1000 # ms

self.timeout = 5000 # ms

self.locale = None # see https://playwright.dev/python/docs/api/class-browser#browser-new-context-option-locale

self.timezone_id = None # see https://playwright.dev/python/docs/api/class-browser#browser-new-context-option-timezone-id

@abstractmethod

def setup(self, page: playwright.sync_api.Page) -> tuple[str, dict]:

"""

Set up everything needed to execute the task.

Args:

page: the active playwright page.

Returns:

goal: str, goal of the task.

info: dict, custom information from the task.

"""

@abstractmethod

def validate(

self, page: playwright.sync_api.Page, chat_messages: list[str]

) -> Tuple[float, bool, str, dict]:

"""

Validate the task was completed successfully

Args:

page: the active playwright page.

chat_messages: the chat messages.

Returns:

reward: float, the reward obtained since last call to validate().

done: boolean flag, indicates if the task has finished or not (be it success or fail).

message: string, a new user message for the chat.

info: dictionnary, custom information from the task.

"""

def cheat(self, page: playwright.sync_api.Page, chat_messages: list[str]) -> None:

"""

Solve the task using a pre-defined solution (optional).

"""

raise NotImplementedError

def teardown(self) -> None:

"""

Tear down the task and clean up any resource / data created by the task (optional).

"""

pass

class OpenEndedTask(AbstractBrowserTask):

@classmethod

def get_task_id(cls):

return "openended"

def __init__(self, seed: int, start_url: str, goal: str = None) -> None:

"""

Args:

seed: random seed.

start_url: str, the url for the starting page.

goal: str, the initial goal.

"""

super().__init__(seed)

self.start_url = start_url

self.goal = goal

def setup(self, page: playwright.sync_api.Page) -> tuple[str, dict]:

page.goto(self.start_url, timeout=10000)

return self.goal, {}

def teardown(self) -> None:

pass

def validate(

self, page: playwright.sync_api.Page, chat_messages: list[str]

) -> Tuple[float, bool, str, dict]:

reward, done, msg, info = 0, False, "", {}

for message in chat_messages:

if message["role"] == "user" and message["message"] == "exit":

done = True

break

return reward, done, msg, info