Jaswin-27/MetaXSCALER_Hackathon

1from models import State, Action, StepResult
2
3
4class LifeEnv:
5
6    def __init__(self):
7        self.state = None
8
9    # RESET ENVIRONMENT
10    def reset(self):
11        self.state = State(
12            energy=100,
13            money=0,
14            happiness=50,
15            day=1
16        )
17        return self.state
18
19    # STEP FUNCTION (MOST IMPORTANT)
20    def step(self, action: Action) -> StepResult:
21
22        if action.type == "work":
23            self.state.money += 50
24            self.state.energy -= 20
25            self.state.happiness -= 5
26
27        elif action.type == "rest":
28            self.state.energy += 30
29            self.state.happiness += 10
30
31        elif action.type == "study":
32            self.state.energy -= 10
33            self.state.happiness -= 2
34            self.state.money += 10  # future benefit
35
36        # increase day
37        self.state.day += 1
38
39        # reward function (IMPORTANT)
40        reward = (
41            self.state.money * 0.1 +
42            self.state.happiness * 0.5 +
43            self.state.energy * 0.2
44        )
45
46        # done condition
47        done = self.state.day > 10 or self.state.energy <= 0
48
49        return StepResult(
50            state=self.state,
51            reward=reward,
52            done=done
53        )
54
55    # GET CURRENT STATE
56    def get_state(self):
57        return self.state