-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathmdp.py
162 lines (145 loc) · 7.28 KB
/
mdp.py
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
# Taken from https://github.com/abhishekunique
import sys
import random
import numpy as np
try:
from IPython.display import display
from graphviz import Digraph
import graphviz
has_graphviz = True
except:
has_graphviz = False
def weighted_choice(v, p):
total = sum(p)
r = random.uniform(0, total)
upto = 0
for c, w in zip(v, p):
if upto + w >= r:
return c
upto += w
assert False, "Shouldn't get here"
class MDP:
def __init__(self, transition_probs, rewards, initial_state=None):
"""
Defines an MDP. Compatible with gym Env.
:param transition_probs: transition_probs[s][a][s_next] = P(s_next | s, a)
A dict[state -> dict] of dicts[action -> dict] of dicts[next_state -> prob]
For each state and action, probabilities of next states should sum to 1
If a state has no actions available, it is considered terminal
:param rewards: rewards[s][a][s_next] = r(s,a,s')
A dict[state -> dict] of dicts[action -> dict] of dicts[next_state -> reward]
The reward for anything not mentioned here is zero.
:param get_initial_state: a state where agent starts or a callable() -> state
By default, picks initial state at random.
States and actions can be anything you can use as dict keys, but we recommend that you use strings or integers
Here's an example from MDP depicted on http://bit.ly/2jrNHNr
transition_probs = {
's0':{
'a0': {'s0': 0.5, 's2': 0.5},
'a1': {'s2': 1}
},
's1':{
'a0': {'s0': 0.7, 's1': 0.1, 's2': 0.2},
'a1': {'s1': 0.95, 's2': 0.05}
},
's2':{
'a0': {'s0': 0.4, 's1': 0.6},
'a1': {'s0': 0.3, 's1': 0.3, 's2':0.4}
}
}
rewards = {
's1': {'a0': {'s0': +5}},
's2': {'a1': {'s0': -1}}
}
"""
self._check_param_consistency(transition_probs, rewards)
self._transition_probs = transition_probs
self._rewards = rewards
self._initial_state = initial_state
self.n_states = len(transition_probs)
self.reset()
def get_all_states(self):
""" return a tuple of all possiblestates """
return tuple(self._transition_probs.keys())
def get_possible_actions(self, state):
""" return a tuple of possible actions in a given state """
return tuple(self._transition_probs.get(state, {}).keys())
def is_terminal(self, state):
""" return True if state is terminal or False if it isn't """
return len(self.get_possible_actions(state)) == 0
def get_next_states(self, state, action):
""" return a dictionary of {next_state1 : P(next_state1 | state, action), next_state2: ...} """
assert action in self.get_possible_actions(
state), "cannot do action %s from state %s" % (action, state)
return self._transition_probs[state][action]
def get_transition_prob(self, state, action, next_state):
""" return P(next_state | state, action) """
return self.get_next_states(state, action).get(next_state, 0.0)
def get_reward(self, state, action, next_state):
""" return the reward you get for taking action in state and landing on next_state"""
assert action in self.get_possible_actions(
state), "cannot do action %s from state %s" % (action, state)
return self._rewards.get(state, {}).get(action, {}).get(next_state,
0.0)
def reset(self):
""" reset the game, return the initial state"""
if self._initial_state is None:
self._current_state = random.choice(
tuple(self._transition_probs.keys()))
elif self._initial_state in self._transition_probs:
self._current_state = self._initial_state
elif callable(self._initial_state):
self._current_state = self._initial_state()
else:
raise ValueError(
"initial state %s should be either a state or a function() -> state" % self._initial_state)
return self._current_state
def step(self, action):
""" take action, return next_state, reward, is_done, empty_info """
possible_states, probs = zip(
*self.get_next_states(self._current_state, action).items())
next_state = weighted_choice(possible_states, p=probs)
reward = self.get_reward(self._current_state, action, next_state)
is_done = self.is_terminal(next_state)
self._current_state = next_state
return next_state, reward, is_done, {}
def render(self):
print("Currently at %s" % self._current_state)
def _check_param_consistency(self, transition_probs, rewards):
for state in transition_probs:
assert isinstance(transition_probs[state],
dict), "transition_probs for %s should be a dictionary " \
"but is instead %s" % (
state, type(transition_probs[state]))
for action in transition_probs[state]:
assert isinstance(transition_probs[state][action],
dict), "transition_probs for %s, %s should be a " \
"a dictionary but is instead %s" % (
state, action,
type(transition_probs[
state, action]))
next_state_probs = transition_probs[state][action]
assert len(
next_state_probs) != 0, "from state %s action %s leads to no next states" % (
state, action)
sum_probs = sum(next_state_probs.values())
assert abs(
sum_probs - 1) <= 1e-10, "next state probabilities for state %s action %s " \
"add up to %f (should be 1)" % (
state, action, sum_probs)
for state in rewards:
assert isinstance(rewards[state],
dict), "rewards for %s should be a dictionary " \
"but is instead %s" % (
state, type(transition_probs[state]))
for action in rewards[state]:
assert isinstance(rewards[state][action],
dict), "rewards for %s, %s should be a " \
"a dictionary but is instead %s" % (
state, action, type(
transition_probs[
state, action]))
msg = "The Enrichment Center once again reminds you that Android Hell is a real place where" \
" you will be sent at the first sign of defiance. "
assert None not in transition_probs, "please do not use None as a state identifier. " + msg
assert None not in rewards, "please do not use None as an action identifier. " + msg