saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-07-23 19:32:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Bug fix for discrete actions

Browse source
This commit is contained in:
Antonin Raffin 2019-10-25 12:00:37 +02:00
parent f8bcb8ee16
commit 584f549fa1
2 changed files with 4 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
torchy_baselines/common/buffers.py
View file

@ -134,6 +134,7 @@ class RolloutBuffer(BaseBuffer):
if len(log_prob.shape) == 0:
# Reshape 0-d tensor to avoid error
log_prob = log_prob.reshape(-1, 1)
self.observations[self.pos] = th.FloatTensor(np.array(obs).copy())
self.actions[self.pos] = th.FloatTensor(np.array(action).copy())
self.rewards[self.pos] = th.FloatTensor(np.array(reward).copy())

3
torchy_baselines/ppo/ppo.py
View file

@ -158,6 +158,9 @@ class PPO(BaseRLModel):
self._update_info_buffer(infos)
n_steps += 1
if isinstance(self.action_space, gym.spaces.Discrete):
# Reshape in case of discrete action
actions = actions.reshape(-1, 1)
rollout_buffer.add(obs, actions, rewards, dones, values, log_probs)
obs = new_obs