def update_policy(self, dataset, epoch=1):
""" Update policy
Args:
paths: a list of trajectories. Each contain a list of symbolic log_prob and rewards
Returns:
"""
actions, advantage, observation, rewards, hidden, mask = dataset
observation = convert_numpy_to_tensor(observation)
actions = convert_numpy_to_tensor(actions)
advantage = convert_numpy_to_tensor(advantage)
rewards = convert_numpy_to_tensor(rewards)
for _ in range(epoch):
# update policy network
self.policy_optimizer.zero_grad()
# compute log prob, assume observation is small.
if not self.recurrent:
distribution, _, raw_baselines = self.policy_net.forward(
observation, None)
log_prob = distribution.log_prob(actions)
else:
log_prob = []
raw_baselines = []
zero_index = np.where(mask == 0)[0] + 1
zero_index = zero_index.tolist()
zero_index.insert(0, 0)
for i in range(len(zero_index) - 1):
start_index = zero_index[i]
end_index = zero_index[i + 1]
current_obs = observation[start_index:end_index]
current_actions = actions[start_index:end_index]
current_hidden = convert_numpy_to_tensor(
np.expand_dims(self.init_hidden_unit, axis=0))
current_dist, _, current_baseline = self.policy_net.forward(
current_obs, current_hidden)
log_prob.append(current_dist.log_prob(current_actions))
raw_baselines.append(current_baseline)
log_prob = torch.cat(log_prob, dim=0)
raw_baselines = torch.cat(raw_baselines, dim=0)
assert log_prob.shape == advantage.shape, 'log_prob length {}, advantage length {}'.format(
log_prob.shape, advantage.shape)
action_loss = torch.mean(-log_prob * advantage)
loss = action_loss
if self.nn_baseline:
value_loss = self.get_baseline_loss(raw_baselines, rewards)
loss = loss + value_loss * self.value_coef
nn.utils.clip_grad_norm_(self.policy_net.parameters(),
self.max_grad_norm)
loss.backward()
self.policy_optimizer.step()
def predict_next_state(self, state, action):
states = np.expand_dims(state, axis=0)
actions = np.expand_dims(action, axis=0)
states = convert_numpy_to_tensor(states)
actions = convert_numpy_to_tensor(actions)
with torch.no_grad():
next_state = self.predict_next_states(states,
actions).cpu().numpy()[0]
return next_state
def predict(self, history_state, history_actions, current_state):
"""
Args:
history_state: (T - 1, 6)
history_actions: (T - 1, 4)
current_state: (6,)
Returns: best action (4,)
"""
states = np.expand_dims(history_state, axis=0) # (1, T - 1, 6)
states = np.tile(
states, (self.num_random_action_selection, 1, 1)) # (N, T - 1, 6)
states = convert_numpy_to_tensor(states)
next_states = np.expand_dims(current_state, axis=0) # (1, 6)
next_states = np.tile(next_states,
(self.num_random_action_selection, 1)) # (N, 6)
next_states = convert_numpy_to_tensor(next_states)
actions = self.action_sampler.sample(
(self.horizon, self.num_random_action_selection)) # (H, N, 4)
actions = convert_numpy_to_tensor(actions)
history_actions = np.expand_dims(history_actions,
axis=0) # (1, T - 1, 4)
current_action = np.tile(
history_actions,
(self.num_random_action_selection, 1, 1)) # (N, T - 1, 4)
current_action = convert_numpy_to_tensor(current_action)
with torch.no_grad():
cost = torch.zeros(
size=(self.num_random_action_selection, )).type(FloatTensor)
for i in range(self.horizon):
states = torch.cat(
(states, torch.unsqueeze(next_states, dim=1)),
dim=1) # # (N, T, 6)
current_action = torch.cat(
(current_action, torch.unsqueeze(actions[i], dim=1)),
dim=1) # (N, T, 4)
next_states = self.model.predict_next_states(
states, current_action) # (N, 6)
cost += self.cost_fn(states[:, -1, :], actions[i],
next_states) * self.gamma_inverse
current_action = current_action[:, 1:, :] # (N, T - 1, 4)
states = states[:, 1:, :]
best_action = actions[0, torch.argmin(cost, dim=0)]
best_action = best_action.cpu().numpy()
return best_action
def predict(self, history_state, history_action, state):
state = np.expand_dims(state, axis=0)
history_state = np.expand_dims(history_state, axis=0)
history_action = np.expand_dims(history_action, axis=0)
with torch.no_grad():
state = convert_numpy_to_tensor(state)
history_state = convert_numpy_to_tensor(history_state)
history_action = convert_numpy_to_tensor(history_action)
state = (state - self.state_mean.squeeze(dim=1)
) / self.state_std.squeeze(dim=1)
history_state = (history_state - self.state_mean) / self.state_std
action = self.model.forward(history_state, history_action, state)
return action.cpu().numpy()[0]
def predict(self, state):
state = np.expand_dims(state, axis=0)
with torch.no_grad():
state = convert_numpy_to_tensor(state)
state = (state - self.state_mean) / self.state_std
action = self.model.forward(state)
return action.cpu().numpy()[0]
def __init__(self, env, temperature_center):
super(EnergyPlusObsWrapper, self).__init__(env=env)
self.obs_mean = np.array([
temperature_center, temperature_center, temperature_center, 1e5,
5000.
],
dtype=np.float32)
self.obs_max = np.array([30., 30., 30., 1e5, 1e4], dtype=np.float32)
self.obs_mean_tensor = convert_numpy_to_tensor(
self.obs_mean).unsqueeze(dim=0)
self.obs_max_tensor = convert_numpy_to_tensor(
self.obs_max).unsqueeze(dim=0)
self.observation_space = spaces.Box(
low=np.array([-1., -1., -1., -10., -10.]),
high=np.array([1., 1., 1., 10.0, 10.0]),
dtype=np.float32)
def predict(self, state):
states = np.expand_dims(state, axis=0)
actions = self.action_sampler.sample(
(self.horizon, self.num_random_action_selection))
states = np.tile(states, (self.num_random_action_selection, 1))
states = convert_numpy_to_tensor(states)
actions = convert_numpy_to_tensor(actions)
with torch.no_grad():
cost = torch.zeros(
size=(self.num_random_action_selection, )).type(FloatTensor)
for i in range(self.horizon):
next_states = self.model.predict_next_states(
states, actions[i])
cost += self.cost_fn(states, actions[i],
next_states) * self.gamma_inverse
states = next_states
best_action = actions[0, torch.argmin(cost, dim=0)]
best_action = best_action.cpu().numpy()
return best_action
def set_statistics(self, dataset):
self.state_mean = convert_numpy_to_tensor(
dataset.state_mean).unsqueeze(dim=0)
self.state_std = convert_numpy_to_tensor(
dataset.state_std).unsqueeze(dim=0)
if self.dynamics_model.discrete:
self.action_mean = None
self.action_std = None
else:
self.action_mean = convert_numpy_to_tensor(
dataset.action_mean).unsqueeze(dim=0)
self.action_std = convert_numpy_to_tensor(
dataset.action_std).unsqueeze(dim=0)
self.delta_state_mean = convert_numpy_to_tensor(
dataset.delta_state_mean).unsqueeze(dim=0)
self.delta_state_std = convert_numpy_to_tensor(
dataset.delta_state_std).unsqueeze(dim=0)
def predict(self, state):
""" The model must be in evaluation mode and turn off gradient update
Args:
state: (ob_dim)
Returns: optimal action (ac_dim)
"""
action_module = TanhActionModule(
init_action=self.action_sampler.sample((self.horizon, )))
optimizer = torch.optim.Adam(action_module.parameters(), lr=1e-3)
# t = tqdm(range(self.num_iterations), desc='Planning')
t = range(self.num_iterations)
for iteration in t:
optimizer.zero_grad()
cost = []
current_state = convert_numpy_to_tensor(
np.expand_dims(state, axis=0))
for h in range(self.horizon):
current_action = action_module.forward(h)
next_states = self.model.predict_next_states(
current_state, current_action)
cost.append(
self.cost_fn(current_state, current_action, next_states) *
self.gamma_inverse)
current_state = next_states
cost = torch.mean(torch.cat(cost))
cost.backward()
nn.utils.clip_grad_norm_(action_module.parameters(), max_norm=1.0)
optimizer.step()
# t.set_description('Iter {}/{}, Cost {:.4f}'.format(iteration + 1, self.num_iterations, cost.item()))
return action_module.forward(0)[0].cpu().detach().numpy()
def __init__(self, init_action):
super(TanhActionModule, self).__init__()
init_action = convert_numpy_to_tensor(init_action)
self.action = nn.Parameter(data=init_action, requires_grad=True)
def set_state_stats(self, state_mean, state_std):
self.state_mean = convert_numpy_to_tensor(state_mean).unsqueeze(dim=0)
self.state_std = convert_numpy_to_tensor(state_std).unsqueeze(dim=0)
def compute_reward_to_go_gae(paths, gamma, policy_net, lam, value_mean,
value_std):
rewards = []
gaes = []
for path in paths:
# compute last state value
if path['mask'][-1] == 1:
with torch.no_grad():
last_obs = convert_numpy_to_tensor(
np.expand_dims(path['last_obs'], axis=0)).type(FloatTensor)
last_hidden = convert_numpy_to_tensor(
np.expand_dims(path['last_hidden'],
axis=0)).type(FloatTensor)
last_state_value = policy_net.forward(
last_obs, last_hidden)[-1].cpu().numpy()[0]
last_state_value = last_state_value * value_std + value_mean
else:
last_state_value = 0.
# we need to clip last_state_value by (max_abs_value / (1 - gamma))
# Otherwise, large state value would cause positive feedback loop and cause the reward to explode.
max_abs_value = np.max(np.abs(path['reward']))
last_state_value = np.clip(last_state_value,
a_min=-max_abs_value / (1 - gamma),
a_max=max_abs_value / (1 - gamma))
# calculate reward-to-go
path['reward'].append(last_state_value)
current_rewards = discount(path['reward'], gamma).astype(np.float32)
rewards.append(current_rewards[:-1])
# compute gae
with torch.no_grad():
observation = path['observation']
hidden = path['hidden']
data_loader = create_data_loader((observation, hidden),
batch_size=32,
shuffle=False,
drop_last=False)
values = []
for obs, hid in data_loader:
obs = move_tensor_to_gpu(obs)
hid = move_tensor_to_gpu(hid)
values.append(policy_net.forward(obs, hid)[-1])
values = torch.cat(values, dim=0).cpu().numpy()
values = values * value_std + value_mean
values = np.append(values, last_state_value)
# add the value of last obs for truncated trajectory
temporal_difference = path[
'reward'][:-1] + values[1:] * gamma - values[:-1]
# calculate reward-to-go
gae = discount(temporal_difference, gamma * lam).astype(np.float32)
gaes.append(gae)
rewards = np.concatenate(rewards)
new_values_mean, new_values_std = np.mean(rewards), np.std(rewards)
rewards = (rewards - new_values_mean) / (new_values_std + eps)
gaes = np.concatenate(gaes)
gaes = (gaes - np.mean(gaes)) / (np.std(gaes) + eps)
return rewards, gaes, new_values_mean, new_values_std