JuliaPOMDP / RockSample.jl

Implement the rock sample problem using POMDPs.jl
Other
5 stars 5 forks source link

Faster State Iteration #18

Closed WhiffleFish closed 2 years ago

WhiffleFish commented 2 years ago

Previous state iterator called state_from_index on each iteration which in turn creates CartesianIndices for which only a single element is used.

In this PR, the temporary CartesianIndices allocation at each iteration is removed by allocating at the very beginning then including it in the iteration state.

The following benchmark is used for comparison:

rs = RockSamplePOMDP(5,7)
@benchmark collect(states(rs))

Previously, we get mean time of: 3.839 ms ± 2.036 ms Now, we get: 731.324 μs ± 643.553 μs

This is a ~5x speedup, and the new array of collected states is equivalent to the vector created with the old iteration method (verified only on RockSamplePOMDP(5,7))

WhiffleFish commented 2 years ago

Benchmark above brought down to 361.415 μs ± 341.099 μs by making rock state static vector creation nonallocating (with @view), and faster "closed-form" solution to CartesianIndices(v)[idx].

WhiffleFish commented 2 years ago

Yeah you're entirely right. My changes were made under the impression that CartesianIndices actually allocated a whole vector filled with CartesianIndex, but that's clearly wrong. One of the problems was, as you stated, splatting the rocks_dim Vector into the CartesianIndices constructor, but the other big issue was the rocks static array:

rocks = SVector{K, Bool}([(s[i] - 1) for i=3:K+2])

The previous code allocated space for a regular base julia vector before converting to static vector. The new code performs this same calculation without the comprehension that was kicking up a lot of garbage collection.

zsunberg commented 2 years ago

Awesome, thanks!!