Differentiable ABMs with StochasticAD.jl #175

kavir1698 · 2020-03-26T16:04:04Z

EDIT: The issue title has been edited to showcase that there is now a formal solution on how to differentiate ABMs, see this comment. Implementing a model that integrates with this as one of our "Integration Examples" would be lovely. :)

Perhaps we can even provide some function(s) that make integration easy, as we did when we were writing the integration example with CellListMap.jl (#659 )

I have tried to use Zygote.jl to differentiate some random and small ABMs. But I get many different errors. Are there data structures that are not supported by Zygote yet? A differentiable ABM would be amazing and allowing to solve many problems.

Datseris · 2020-03-26T16:22:10Z

I know nothing about zygote, nor the data structures it needs. I can't add anything to the discussion here. :(

p.s.: This issue is not well described. You should consider adding more information about the errors you get, what you are trying to do, a MWE, and more.

kavir1698 · 2020-03-26T16:45:49Z

Agent-based model are often very complex and one needs to run many simulations and replicates to explore parameter space and understand the behavior of a model. If we were able to differentiate AMBs with Automatic Differentiation (AD), we would be able to save a lot of time in finding parameter values that best fit data.

I have read about AD in Julia, but I am not deeply familiar with its limitations. I tried to differentiate a simple ABM (Agent.jl's wealth distribution example). But I get the following error. I am not sure what part of the data structure we use are not supported.

using Agents
using Zygote
using Statistics: std

mutable struct WealthAgent <: AbstractAgent
    id::Int
    wealth::Int
end

function wealth_model(;numagents = 100, initwealth = 1)
    model = ABM(WealthAgent, scheduler=random_activation)
    for i in 1:numagents
        add_agent!(model, initwealth)
    end
    return model
end

function agent_step!(agent, model)
    agent.wealth == 0 && return # do nothing
    ragent = random_agent(model)
    agent.wealth -= 1
    ragent.wealth += 1
end

function costWealth(M)
  agent_properties = [:wealth]
  model = wealth_model(numagents=M)
  step!(model, agent_step!, 5)
  std([agent.wealth for agent in values(model.agents)])
end

gradient(N ->costWealth(N), 1000)

ERROR: InexactError: Int64(0.002002002002002002)
Stacktrace:
 [1] Int64 at .\float.jl:709 [inlined]
 [2] convert at .\number.jl:7 [inlined]
 [3] _backvar at C:\Users\Ali\.julia\packages\Zygote\KNUTW\src\lib\array.jl:270 
[inlined]
 [4] _backvar at C:\Users\Ali\.julia\packages\Zygote\KNUTW\src\lib\array.jl:269 
[inlined]
 [5] #1264 at C:\Users\Ali\.julia\packages\Zygote\KNUTW\src\lib\array.jl:274 [inlined]
 [6] (::Zygote.var"#3219#back#1266"{Zygote.var"#1264#1265"{Bool,Colon,Float64,Array{Int64,1},Float64}})(::Float64) at C:\Users\Ali\.julia\packages\ZygoteRules\6nssF\src\adjoint.jl:49
 [7] costWealth at .\REPL[115]:7 [inlined]
 [8] (::typeof(∂(costWealth)))(::Float64) at C:\Users\Ali\.julia\packages\Zygote\KNUTW\src\compiler\interface2.jl:0
 [9] #17 at .\REPL[119]:1 [inlined]
 [10] (::Zygote.var"#38#39"{typeof(∂(#17))})(::Float64) at C:\Users\Ali\.julia\packages\Zygote\KNUTW\src\compiler\interface.jl:36
 [11] gradient(::Function, ::Int64, ::Vararg{Int64,N} where N) at C:\Users\Ali\.julia\packages\Zygote\KNUTW\src\compiler\interface.jl:45
 [12] top-level scope at REPL[119]:1

The error varies depending of the model though.

Libbum · 2020-03-26T21:51:06Z

I'm getting a core dump when I try this example

julia: /buildworker/worker/package_linux64/build/src/codegen.cpp:4357: jl_cgval_t emit_expr(jl_codectx_t&, jl_value_t*, ssize_t): Assertion `token.V->getType()->isTokenTy()' failed.
so I can't help troubleshoot until I figure that out.

I haven't worked with AD in Julia that much myself, but I've never considered using it on discrete variables. What if you changed wealth to Float64?

kavir1698 · 2020-03-26T22:52:08Z

Thanks for the reply. Good point. I get the error below with Float64 for wealth. But this error seems more promising to be solvable.

ERROR: Need an adjoint for constructor Base.ValueIterator{Dict{Int64,WealthAgent}}. Gradient is of type Array{Base.RefValue{Any},1}
Stacktrace:
 [1] error(::String) at .\error.jl:33
 [2] (::Zygote.Jnew{Base.ValueIterator{Dict{Int64,WealthAgent}},Nothing,false})(::Array{Base.RefValue{Any},1}) at C:\Users\Ali\.julia\packages\Zygote\KNUTW\src\lib\lib.jl:294
 [3] (::Zygote.var"#378#back#196"{Zygote.Jnew{Base.ValueIterator{Dict{Int64,WealthAgent}},Nothing,false}})(::Array{Base.RefValue{Any},1}) at C:\Users\Ali\.julia\packages\ZygoteRules\6nssF\src\adjoint.jl:49
 [4] ValueIterator at .\abstractdict.jl:44 [inlined]
 [5] (::typeof(∂(Base.ValueIterator{Dict{Int64,WealthAgent}})))(::Array{Base.RefValue{Any},1}) at C:\Users\Ali\.julia\packages\Zygote\KNUTW\src\compiler\interface2.jl:0
 [6] ValueIterator at .\abstractdict.jl:44 [inlined]
 [7] (::typeof(∂(Base.ValueIterator)))(::Array{Base.RefValue{Any},1}) at C:\Users\Ali\.julia\packages\Zygote\KNUTW\src\compiler\interface2.jl:0
 [8] values at .\abstractdict.jl:123 [inlined]
 [9] (::typeof(∂(values)))(::Array{Base.RefValue{Any},1}) at C:\Users\Ali\.julia\packages\Zygote\KNUTW\src\compiler\interface2.jl:0
 [10] costWealth at .\REPL[8]:5 [inlined]
 [11] (::typeof(∂(costWealth)))(::Float64) at C:\Users\Ali\.julia\packages\Zygote\KNUTW\src\compiler\interface2.jl:0
 [12] #6 at .\REPL[10]:1 [inlined]
 [13] (::Zygote.var"#38#39"{typeof(∂(#6))})(::Float64) at C:\Users\Ali\.julia\packages\Zygote\KNUTW\src\compiler\interface.jl:36
 [14] gradient(::Function, ::Int64, ::Vararg{Int64,N} where N) at C:\Users\Ali\.julia\packages\Zygote\KNUTW\src\compiler\interface.jl:45
 [15] top-level scope at REPL[10]:1

Libbum · 2020-03-29T19:56:50Z

I'm going to try and dig into this a little deeper. Completely agree it would be good to have this working. Luckily I've had a need to start working with Zygote myself recently, so hopefully we can get this done for v3

ChrisRackauckas · 2020-04-03T00:44:32Z

You could attempt to get the strong solution of some models, but in many ABMs the discreteness is inherently non-differentiable through AD, and you'd have to move to the continuous master equation PDE to get a differentiable form.

kavir1698 · 2020-04-03T07:35:45Z

Thank you, @ChrisRackauckas , for your reply. It is great to have the input of an expert on this issue. Is is possible to know when the discreteness is not a problem?
And if not through AD, could numerical differentiation be applied to ABMs?

nunezmatias · 2022-08-06T17:53:16Z

this one might be relevant to this thread "Differentiable Agent-Based Simulation for
Gradient-Guided Simulation-Based Optimization" https://arxiv.org/pdf/2103.12476.pdf

ChrisRackauckas · 2022-08-06T19:44:02Z

That isn't a very great way to do it. That smoothing can introduce a lot of bias and has exponential cost scaling.

We have a paper that will be out really soon that shows how to do this, along with a new AD in Julia that implements the method. I was wrong two years ago that it isn't possible: you just need an entirely different definition of automatic differentiation.

ayushchopra96 · 2022-10-11T13:31:41Z

@ChrisRackauckas this sounds very interesting. Is your paper out yet?
keen to see how you work around sampling from discrete distributions w/o reparameterizing.
Only other way that comes to mind is policy gradients, which isn't really a true gradient..

ChrisRackauckas · 2022-10-11T13:34:25Z

Remind me tomorrow. The library should be out by the end of the week.

arnauqb · 2022-10-12T11:26:11Z

Remind me tomorrow. The library should be out by the end of the week.

Reminder :)

ChrisRackauckas · 2022-10-18T11:39:24Z

Here it is: https://twitter.com/ChrisRackauckas/status/1582324437285625857

https://arxiv.org/abs/2210.08572

Reopen the issue 😄

Datseris · 2022-10-18T13:24:47Z

Implementing a model that integrates with this as one of our "Integration Examples" would be lovely. :)

Perhaps we can even provide some function(s) that make integration easy, as we did when we were writing the integration example with CellListMap.jl (#659 )

ayushchopra96 · 2022-10-18T17:09:48Z

Interesting @ChrisRackauckas. We did similar work with the reparameterization trick to enable differentiable ABMs. Then, we also demonstrate how you can calibrate such large differentiable ABMs with deep neural networks (~10 million agents). https://arxiv.org/abs/2207.09714. Also, at MIT - I would love to chat further and share notes!

Datseris added the help wanted Help from someone external is needed label Mar 26, 2020

kavir1698 closed this as completed Jun 9, 2020

Datseris reopened this Oct 18, 2022

Datseris pinned this issue Oct 18, 2022

Datseris changed the title ~~Differentiable ABM?~~ Differentiable ABMs with StochasticAD.jl Oct 18, 2022

Datseris added documentation example-integration A new example or showcase of integration with the Julia ecosystem labels Oct 18, 2022

Tortar unpinned this issue Jul 12, 2023

Tortar pinned this issue Jul 12, 2023

Tortar unpinned this issue Aug 17, 2023

Tortar pinned this issue Aug 17, 2023

Tortar unpinned this issue Sep 23, 2023

Tortar pinned this issue Feb 29, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Differentiable ABMs with StochasticAD.jl #175

Differentiable ABMs with StochasticAD.jl #175

kavir1698 commented Mar 26, 2020 •

edited by Datseris

Datseris commented Mar 26, 2020

kavir1698 commented Mar 26, 2020 •

edited

Libbum commented Mar 26, 2020

kavir1698 commented Mar 26, 2020

Libbum commented Mar 29, 2020

ChrisRackauckas commented Apr 3, 2020

kavir1698 commented Apr 3, 2020

nunezmatias commented Aug 6, 2022 •

edited

ChrisRackauckas commented Aug 6, 2022 •

edited

ayushchopra96 commented Oct 11, 2022

ChrisRackauckas commented Oct 11, 2022

arnauqb commented Oct 12, 2022

ChrisRackauckas commented Oct 18, 2022 •

edited

Datseris commented Oct 18, 2022

ayushchopra96 commented Oct 18, 2022

Differentiable ABMs with StochasticAD.jl #175

Differentiable ABMs with StochasticAD.jl #175

Comments

kavir1698 commented Mar 26, 2020 • edited by Datseris

Datseris commented Mar 26, 2020

kavir1698 commented Mar 26, 2020 • edited

Libbum commented Mar 26, 2020

kavir1698 commented Mar 26, 2020

Libbum commented Mar 29, 2020

ChrisRackauckas commented Apr 3, 2020

kavir1698 commented Apr 3, 2020

nunezmatias commented Aug 6, 2022 • edited

ChrisRackauckas commented Aug 6, 2022 • edited

ayushchopra96 commented Oct 11, 2022

ChrisRackauckas commented Oct 11, 2022

arnauqb commented Oct 12, 2022

ChrisRackauckas commented Oct 18, 2022 • edited

Datseris commented Oct 18, 2022

ayushchopra96 commented Oct 18, 2022

kavir1698 commented Mar 26, 2020 •

edited by Datseris

kavir1698 commented Mar 26, 2020 •

edited

nunezmatias commented Aug 6, 2022 •

edited

ChrisRackauckas commented Aug 6, 2022 •

edited

ChrisRackauckas commented Oct 18, 2022 •

edited