Troubleshooting jit-ing with methods.

[joseph-jnl]

Hello, I'm going through the Sutton & Barto rlbook while also learning jax (probably too many degrees of freedom to learn efficiently!) and ran into an issue with troubleshooting jit-ing methods.

Specifically, the method works for:

• the first instantiated object
• an object initialized with a different set of parameters

But fails for a second object with the same set of initialized parameters.

I'm unsure if the root cause is incorrectly:

• Defining a custom pytree class. I believe I have all of the arrays as children and the other attributes as static:

    def _tree_flatten(self):
        children = (
            self.R,
            self.P,
            self.actions,
            self.v_init,
        )  # arrays and dynamic values
        # static values (non-arrays)
        aux_data = {
            "special_states": self.special_states,
            "special_states_prime": self.special_states_prime,
            "special_states_rewards": self.special_states_rewards,
        }

        return (children, aux_data)

    @classmethod
    def _tree_unflatten(cls, aux_data, children):
        grid = cls(
            aux_data["special_states"],
            aux_data["special_states_prime"],
            aux_data["special_states_rewards"],
            R=children[0],
            P=children[1],
        )
        grid.v_init = children[2]

        return grid

• Defining a pure function, I believe I've removed all references to state:

    def state_value(
        self,
        v,
        R,
        P,
        special_states,
        special_states_prime,
        special_states_rewards,
        discount: float = 0.9,
    ):
        """"""
        # Update interior grid
        vp = (
            R
            + convolve2d(
                jnp.pad(v, pad_width=(1, 1), constant_values=0),
                P,
                mode="valid",
            )
            * discount
        )

        # Update edges except for corners
        vp = vp.at[1:-1, 0].add(v[1:-1, 0] * discount * 0.25)
        vp = vp.at[1:-1, -1].add(v[1:-1, -1] * discount * 0.25)
        vp = vp.at[0, 1:-1].add(v[0, 1:-1] * discount * 0.25)
        vp = vp.at[-1, 1:-1].add(v[-1, 1:-1] * discount * 0.25)

        # Update corners
        vp = vp.at[0, 0].add(v[0, 0] * 2 * discount * 0.25)
        vp = vp.at[0, -1].add(v[0, -1] * 2 * discount * 0.25)
        vp = vp.at[-1, -1].add(v[-1, -1] * 2 * discount * 0.25)
        vp = vp.at[-1, 0].add(v[-1, 0] * 2 * discount * 0.25)

        # Update special states
        vp = vp.at[special_states[0], special_states[1]].set(
            v[special_states_prime[0], special_states_prime[1]] * discount
            + special_states_rewards
        )

        return vp

• Calling the function in the final loop- I'm uncertain on whether the final calling function can also be a method?:

    def estimate_state_value(self, iter=1000):
        """"""
        v = self.v_init
        for _ in range(iter):
            v = self.state_value(
                v,
                self.R,
                self.P,
                self.special_states,
                self.special_states_prime,
                self.special_states_rewards,
            )
        return v

• Something else I haven't thought of

Link to full code.
Thank you for any thoughts!

[jakevdp]

The issue is that your tree_flatten method includes the array special_states_rewards in aux_data, and arrays are not allowed in aux_data because they are not hashable and don't have simple equality semantics. From https://jax.readthedocs.io/en/latest/pytrees.html:

When defining unflattening functions, in general children should contain all the dynamic elements of the data structure (arrays, dynamic scalars, and pytrees), while aux_data should contain all the static elements that will be rolled into the treedef structure. JAX sometimes needs to compare treedef for equality, or compute its hash for use in the JIT cache, and so care must be taken to ensure that the auxiliary data specified in the flattening recipe supports meaningful hashing and equality comparisons.

The fix would be to include special_states_rewards in the children within the flattening rule.

[joseph-jnl]

I should probably read my own type hints! Thank you!