From 04ecd289701a2273366c0ffc1d4a1f5f60d99cba Mon Sep 17 00:00:00 2001 From: Graeme22 Date: Thu, 8 Jun 2023 10:57:10 -0500 Subject: [PATCH] fix bug in masked, multi-dimensional state space multi-threaded envs --- .../src/algorithms/policy_gradient/multi_thread_env.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ReinforcementLearningZoo/src/algorithms/policy_gradient/multi_thread_env.jl b/src/ReinforcementLearningZoo/src/algorithms/policy_gradient/multi_thread_env.jl index 11d175833..10a17157c 100644 --- a/src/ReinforcementLearningZoo/src/algorithms/policy_gradient/multi_thread_env.jl +++ b/src/ReinforcementLearningZoo/src/algorithms/policy_gradient/multi_thread_env.jl @@ -131,7 +131,7 @@ function RLBase.is_terminated(env::MultiThreadEnv) end function RLBase.legal_action_space_mask(env::MultiThreadEnv) - N = ndims(env.states) + N = ndims(env.legal_action_space_mask) @sync for i in 1:length(env) @spawn selectdim(env.legal_action_space_mask, N, i) .= legal_action_space_mask(env[i])