Skip to content

Commit

Permalink
Simplify FluxApproximator's optimise! method by using a single-line f…
Browse files Browse the repository at this point in the history
…unction definition
  • Loading branch information
jeremiahpslewis committed Dec 17, 2024
1 parent de3652d commit b0e7e85
Showing 1 changed file with 1 addition and 4 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -43,8 +43,5 @@ Flux.@layer FluxApproximator trainable=(model,)
forward(A::FluxApproximator, args...; kwargs...) = A.model(args...; kwargs...)
forward(A::FluxApproximator, env::E, player::AbstractPlayer=current_player(env)) where {E <: AbstractEnv} = env |> (x -> state(x, player)) |> (x -> forward(A, x))

function RLBase.optimise!(A::FluxApproximator, grad::NamedTuple)

Flux.Optimise.update!(A.optimiser_state, A.model, grad.model)
end
RLBase.optimise!(A::FluxApproximator, grad::NamedTuple) = Flux.Optimise.update!(A.optimiser_state, A.model, grad.model)

0 comments on commit b0e7e85

Please sign in to comment.