[content] genie.py (python)
    def reset_spatial_kv_cache_direct(self):
        """
        Directly resets the spatial key-value cache for all STBlocks in the dynamics model
        by accessing the cache variables without using the wrapper methods.
        """
        for i in range(self.dynamics.dynamics.num_blocks):
            st_block = getattr(self.dynamics.dynamics, f'STBlock_{i}')
            # In STBlock, the first MultiHeadAttention is spatial (index 0)
            spatial_attn = getattr(st_block, 'MultiHeadAttention_0')
            if spatial_attn.has_variable('cache', 'cached_key'):
                cached_key = spatial_attn.get_variable('cache', 'cached_key')
                cached_value = spatial_attn.get_variable('cache', 'cached_value')
                
                spatial_attn.put_variable('cache', 'cached_key', jnp.zeros_like(cached_key))
                spatial_attn.put_variable('cache', 'cached_value', jnp.zeros_like(cached_value))
                spatial_attn.put_variable('cache', 'cache_index', jnp.array(0, dtype=jnp.int32))

    def reset_spatial_kv_cache_most_direct(self):
        """
        The most direct way to reset spatial KV cache - accessing cache variables
        directly from the MultiHeadAttention modules without any wrapper methods.
        This shows the exact cache variable names and structure.
        """
        # Access the dynamics STTransformer
        dynamics_transformer = self.dynamics.dynamics
        
        # Iterate through all STBlocks
        for i in range(dynamics_transformer.num_blocks):
            # Get the STBlock by default name
            st_block = getattr(dynamics_transformer, f'STBlock_{i}')
            
            # Access the spatial attention module directly (first MultiHeadAttention)
            spatial_attn = getattr(st_block, 'MultiHeadAttention_0')
            
            # Check if cache variables exist
            if spatial_attn.has_variable('cache', 'cached_key'):
                # Get current cache values
                cached_key = spatial_attn.get_variable('cache', 'cached_key')
                cached_value = spatial_attn.get_variable('cache', 'cached_value')
                cache_index = spatial_attn.get_variable('cache', 'cache_index')
                
                # Reset to zero/initial state
                spatial_attn.put_variable('cache', 'cached_key', jnp.zeros_like(cached_key))
                spatial_attn.put_variable('cache', 'cached_value', jnp.zeros_like(cached_value))
                spatial_attn.put_variable('cache', 'cache_index', jnp.array(0, dtype=jnp.int32))

    def reset_all_kv_cache_direct(self):
        """
        Demonstrates the most direct way to reset KV cache for both spatial and temporal attention.
        This shows how to access cache variables directly using Flax's default module names.
        """
        dynamics_transformer = self.dynamics.dynamics
        
        for i in range(dynamics_transformer.num_blocks):
            st_block = getattr(dynamics_transformer, f'STBlock_{i}')
            
            # Reset spatial attention cache (first MultiHeadAttention)
            spatial_attn = getattr(st_block, 'MultiHeadAttention_0')
            if spatial_attn.has_variable('cache', 'cached_key'):
                cached_key = spatial_attn.get_variable('cache', 'cached_key')
                cached_value = spatial_attn.get_variable('cache', 'cached_value')
                
                spatial_attn.put_variable('cache', 'cached_key', jnp.zeros_like(cached_key))
                spatial_attn.put_variable('cache', 'cached_value', jnp.zeros_like(cached_value))
                spatial_attn.put_variable('cache', 'cache_index', jnp.array(0, dtype=jnp.int32))
            
            # Reset temporal attention cache (second MultiHeadAttention)
            temporal_attn = getattr(st_block, 'MultiHeadAttention_1')
            if temporal_attn.has_variable('cache', 'cached_key'):
                cached_key = temporal_attn.get_variable('cache', 'cached_key')
                cached_value = temporal_attn.get_variable('cache', 'cached_value')
                
                temporal_attn.put_variable('cache', 'cached_key', jnp.zeros_like(cached_key))
                temporal_attn.put_variable('cache', 'cached_value', jnp.zeros_like(cached_value))
                temporal_attn.put_variable('cache', 'cache_index', jnp.array(0, dtype=jnp.int32))