DefaultTrajectoryState <= GraphTrajectoryState (#428)

corail-research · Jul 13, 2023 · 352d022 · 352d022
1 parent 546ffe3
commit 352d022
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/CP/valueselection/learning/environment.jl b/src/CP/valueselection/learning/environment.jl
@@ -68,7 +68,7 @@ function RLBase.prob(p::PPOPolicy, env::AbstractCPEnv{ST}) where {ST <: NonTabul
     prob(p, s, mask)
 end
 
-function RLBase.prob(p::PPOPolicy{<:ActorCritic,Categorical}, state::DefaultTrajectoryState, mask)
+function RLBase.prob(p::PPOPolicy{<:ActorCritic,Categorical}, state::GraphTrajectoryState, mask)
     logits = p.approximator.actor(send_to_device(device(p.approximator), state))
     if !isnothing(mask)
         logits .+= ifelse.(mask, 0f0, typemin(Float32))