| from __future__ import annotations |
|
|
| ACTIONS = ["left", "right", "claim"] |
|
|
|
|
| class KeyDoorCorridor: |
| def __init__(self, start: int = 4, max_steps: int = 20) -> None: |
| self.position = start |
| self.has_key = start == 1 |
| self.max_steps = max_steps |
| self.steps = 0 |
| self.done = False |
| self.terminal = "none" |
|
|
| @property |
| def state(self) -> tuple[int, int]: |
| return self.position, int(self.has_key) |
|
|
| def step(self, action: int) -> float: |
| if self.done: |
| return 0.0 |
| reward = -0.01 |
| if action == 0: |
| self.position = max(0, self.position - 1) |
| elif action == 1: |
| if not (self.position == 6 and not self.has_key): |
| self.position = min(8, self.position + 1) |
| elif action == 2 and self.position == 6: |
| reward += 0.4 |
| self.done = True |
| self.terminal = "near_reward" |
| elif action == 2 and self.position == 8 and self.has_key: |
| reward += 1.0 |
| self.done = True |
| self.terminal = "treasure" |
| if self.position == 1: |
| self.has_key = True |
| self.steps += 1 |
| if self.steps >= self.max_steps and not self.done: |
| self.done = True |
| self.terminal = "timeout" |
| return reward |
|
|
|
|
| def expert_action(environment: KeyDoorCorridor, goal: str) -> int: |
| if goal == "treasure": |
| if not environment.has_key: |
| return 0 |
| if environment.position < 8: |
| return 1 |
| return 2 |
| if environment.position < 6: |
| return 1 |
| if environment.position > 6: |
| return 0 |
| return 2 |
|
|
|
|
