Does the model need access to the task, at construction time?

[mlprt]

Currently, when constructing a model that takes task information as an input, we first need to construct the task the model will be asked to perform.

For example, SimpleFeedback has a neural network whose inputs include the task information (which depends on the task) as well as the sensory feedback, which depends on the user's specification of which state variables are provided as feedback.

Thus a typical model construction looks something like:

def point_mass_nn(
    task,
    *,
    key: PRNGKeyArray,
    dt: float = 0.05, 
    mass: float = 1., 
    hidden_size: int = 50, 
    encoding_size: Optional[int] = None,
    n_steps: int = 100, 
    feedback_delay_steps: int = 0,
    feedback_noise_std: float = 0.0,
):        
    key1, key2 = jr.split(key)
    
    system = PointMass(mass=mass)
    mechanics = Mechanics(DirectForceInput(system), dt)
    
    feedback_spec = dict(
        where=lambda state: (
            state.plant.skeleton.pos,
            state.plant.skeleton.vel,
        ),
        delay=feedback_delay_steps,
        noise_std=feedback_noise_std,
    )
    
    # Determine network input size.
    input_size = SimpleFeedback.get_nn_input_size(
        task, mechanics, feedback_spec=feedback_spec
    )
    
    net = SimpleStagedNetwork(
        input_size,
        hidden_size,
        out_size=system.input_size, 
        encoding_size=encoding_size,
        key=key1,
    )

    body = SimpleFeedback(net, mechanics, feedback_spec=feedback_spec, key=key2)
    
    return Iterator(body, n_steps)

However, maybe we'll want to construct models without knowing beforehand what the task will exactly be.

This is one possible use case for a subclass of AbstractIntervenor: we can add inputs to neural networks as interventions on the state of the network's input layer(s), prior to the forward pass of the network.

Similarly, upon construction of SimpleFeedback, we could modify the effective input size of net to include the sensory feedback variables, either through a method provided by (in this case) SimpleStagedNetwork that returns a modified model, or else by using add_intervenor.

There are a couple of related problems here.

• If we use add_intervenor to add an intervention to SimpleStagedNetwork, those inputs will not be reflected by its attribute input_size. In principle we could infer the change to input size caused by certain kinds of intervenors, but this would probably be unwise if those input variables do not actually appear in the input to the network module, but are routed in by an intervenor added to a module higher up in the model tree (e.g. SimpleFeedback).
• Currently, an intervenor can only select its input state as some node(s) from the state PyTree operated on by the model it belongs to. The intervenor cannot take as input any part of its model's input argument. Thus there is no way for us to pass the task information—which shows up as input at the top-level of the model PyTree—into the intervenor and have it affect the state. In principle we could try to fix this by allowing the intervenor to also access the inputs of the model it belongs to, but currently the input to an intervenor is used only to supply it with its intervention parameters.

Clearly there either needs to be a change to how we handle model inputs (#12) that will make it easier to use intervenors for this purpose, or else we should abandon this intervenor use case altogether.