Skip to content

Commit

Permalink
Removed AxisArrays in favor of ComponentArrays for significantly redu…
Browse files Browse the repository at this point in the history
…ced first-run latency. Note that accessing named versions of jlModel and jlData may incur slight performance hit if used as naive getindex(q, :root); escaping the name Symbol with Val, i.e. getindex(q, Val(:root)) informs the compiler how to better optimize.
  • Loading branch information
klowrey committed Sep 23, 2021
1 parent af09713 commit 91b2184
Show file tree
Hide file tree
Showing 7 changed files with 27 additions and 18 deletions.
4 changes: 2 additions & 2 deletions Project.toml
Original file line number Diff line number Diff line change
@@ -1,13 +1,13 @@
name = "MuJoCo"
uuid = "93189219-7048-461c-94ec-443a161ed927"
authors = ["Colin Summers", "Kendall Lowrey"]
version = "0.3.1"
version = "0.3.2"

[deps]
AxisArrays = "39de3d68-74b9-583c-8d2d-e117c070f3a9"
BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"
Base64 = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
CEnum = "fa961155-64e5-5f13-b03f-caf6b980ea82"
ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
MuJoCo_jll = "32af7c3b-80ec-5621-8194-2f6cb2280831"
Expand Down
10 changes: 7 additions & 3 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@ All functions are documented for convenience:
help?> mj_step
```

For more complete about the MuJoCo physics simulator, see [MuJoCo's documentation](http://www.mujoco.org/book).
For more complete info about the MuJoCo physics simulator, see [MuJoCo's documentation](http://www.mujoco.org/book).


### MuJoCo
Expand All @@ -67,6 +67,10 @@ m = jlModel("humanoid.xml")
@assert m.opt.timestep == 0.002
@set!! m.opt.timestep = 0.001
@assert m.opt.timestep == 0.001
d = jlData(m)
mn, dn = namify(m, d) # Produced NamedTuples with struct fields and ComponentArrays
dn.qpos[:root_x] # allows users to access MuJoCo fields with names specified in the XML.
dn.qpos[Val(:root_x)] # For better performance, create a Val of the symbol name.
```

#### Globals
Expand Down Expand Up @@ -99,8 +103,8 @@ you will need to add `Lyceum/LyceumRegistry`.
From the Julia REPL, type `]` to enter Pkg mode:
```julia-repl
julia> ]
(v1.3) pkg> registry add https://github.com/Lyceum/LyceumRegistry.git
(v1.3) pkg> add MuJoCo
(v1.6) pkg> registry add https://github.com/Lyceum/LyceumRegistry.git
(v1.6) pkg> add MuJoCo
```

Below we simulate passive dynamics and print out joint positions
Expand Down
2 changes: 1 addition & 1 deletion src/MuJoCo.jl
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,9 @@ module MuJoCo

using MacroTools: combinestructdef
using Base: RefValue
using AxisArrays: Axis, AxisArray
using
UnsafeArrays,
ComponentArrays,
StaticArrays,
Reexport,
BangBang
Expand Down
2 changes: 1 addition & 1 deletion src/jldata.jl
Original file line number Diff line number Diff line change
Expand Up @@ -118,4 +118,4 @@ function build_jlData()
end
end

@eval $(build_jlData())
@eval $(build_jlData())
2 changes: 1 addition & 1 deletion src/metadata.jl
Original file line number Diff line number Diff line change
Expand Up @@ -431,4 +431,4 @@ const STRUCT_INFO = Dict{Symbol,Dict{Symbol,<:FieldInfo}}(
:actuator_length => FieldInfo((:nu,)),
:efc_R => FieldInfo((:njmax,)),
),
)
)
24 changes: 14 additions & 10 deletions src/namify.jl
Original file line number Diff line number Diff line change
Expand Up @@ -219,25 +219,28 @@ end

function namify_array(A, newaxes)
axisarray_axes = map(newaxes) do newax
axname = newax.name
if newax isa UnnamedAxis
Axis{axname}(1:length(newax))
l = length(newax)
if newax isa LinearAxis
keys = Tuple(newax.labels)
vals = Tuple(1:l)
else
Axis{axname}(SVector{length(newax),Symbol}(newax.labels))
keys = (newax.name,)
vals = (1:l,)
end
nt = (; zip(keys, vals)...)
ComponentArrays.Axis(nt)
end
newsize = length.(axisarray_axes)
AxisArray(reshape(A, newsize), axisarray_axes)
newsize = map(length, newaxes)
ComponentArray{axisarray_axes}(reshape(A, newsize))
end

function namify_one(x::Union{jlModel,jlData}, axismapping)
newfields = []
fieldnames = Symbol[]
for fieldname in propertynames(x)
for (fieldname, axis) in axismapping
field = getproperty(x, fieldname)
if field isa Array && haskey(axismapping, fieldname)
newaxes = axismapping[fieldname]
push!(newfields, namify_array(field, newaxes))
if field isa Array
push!(newfields, namify_array(field, axis))
push!(fieldnames, fieldname)
end
end
Expand All @@ -250,3 +253,4 @@ function namify(m::jlModel, d::jlData)
m_axmappings = associate_axes(m, m, sz2ax)
namify_one(m, m_axmappings), namify_one(d, d_axmappings)
end

1 change: 1 addition & 0 deletions test/runtests.jl
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ end

function test_uninit(T)
zeroarg = T()
zeroarg = T() # huh? first empty init may be different
posonly = T((getfield(zeroarg, name) for name in fieldnames(T))...)
isequal(zeroarg, posonly) && hash(zeroarg) == hash(posonly)
end
Expand Down

0 comments on commit 91b2184

Please sign in to comment.