dodox/metastable-baselines2
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Add support for VecEnvs

Browse Source
This commit is contained in:
Dominik Moritz Roth 2024-03-09 12:33:20 +01:00
parent 6e79fce9ae
commit b2384e183c
1 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
metastable_baselines2/common/buffers.py
View File

@ -202,6 +202,8 @@ class BetterRolloutBuffer(RolloutBuffer):
"log_probs", "log_probs",
"advantages", "advantages",
"returns", "returns",
"means",
"cov_decomps"
] ]
for tensor in _tensor_names: for tensor in _tensor_names:
@ -227,8 +229,8 @@ class BetterRolloutBuffer(RolloutBuffer):
self.actions[batch_inds], self.actions[batch_inds],
self.values[batch_inds].flatten(), self.values[batch_inds].flatten(),
self.log_probs[batch_inds].flatten(), self.log_probs[batch_inds].flatten(),
np.squeeze(self.means[batch_inds], axis=1), self.means[batch_inds],
np.squeeze(self.cov_decomps[batch_inds], axis=1), self.cov_decomps[batch_inds],
self.advantages[batch_inds].flatten(), self.advantages[batch_inds].flatten(),
self.returns[batch_inds].flatten(), self.returns[batch_inds].flatten(),
) )