Optimize pickled representation through a custom __getstate__

[cgc]

The pickled representation of a TabularMDP instance can get pretty hefty. We might want to consider a strategy where we have a custom __getstate__ to filter out cached properties (maybe like this blog post?). Might be worth refactoring or cleaning up method_cache to make sure it interoperates well with this.

[markkho]

Are there any cases where we would want that information saved? Seems possible but could get complicated, so I think it makes sense to filter them out.

[cgc]

Definitely agree that it could depend on the usecase!

Since pickling is a handy way to serialize arbitrary objects it winds up being used by several of the caching libraries i’ve been trying. For my case, I basically use the MDP as a key to cache the result of some optimization. When these cached fields are included, the serialized size increases from ~100 bytes to ~7k bytes. It winds up being a cpu bottleneck for me, so I’ve started removing these fields.

One way we could be flexible to usecase is having a flag that controls the behavior in getstate? Something like _pickle_cache or _state_includes_cache.

[cgc]

Have just realized that __getstate__ also seems to be used when running copy.copy(). I do actually use copy.copy() somewhat extensively for a project where I work with MDPs that share a transition function, but with different is_terminal() function. I've come to implicitly rely on the fact that copy.copy() brings cached properties over and usually prime the cache before running my analyses. So, not quite sure what the best approach for this is; may ultimately depend on any refactoring of caching in TabularMDP.

Here's a test case showing the copy.copy behavior for Python 3.8.10:

import copy, pickle

class TestCls:
    def __init__(self):
        self.included_in_getstate = 456
        self.not_included_in_getstate = 789
    def __getstate__(self):
        return dict(included_in_getstate=self.included_in_getstate)

print('object dict', TestCls().__dict__)
# > object dict {'included_in_getstate': 456, 'not_included_in_getstate': 789}
print('after pickling', pickle.loads(pickle.dumps(TestCls())).__dict__)
# > after pickling {'included_in_getstate': 456}
print('after copy', copy.copy(TestCls()).__dict__)
# > after copy {'included_in_getstate': 456}