Clarify Trainer/Algorithm API

[doxav]

Hello,

I’ve been studying the trainer implementation and want to confirm my understanding against the real code and explore possible options.

Current Behavior

1. Agent must be a trace.model (a Module)
   The AlgorithmBase.__init__ asserts the passed agent is a Module—no plain functions or callbacks are accepted .
2. Only single-argument __call__
MinibatchAlgorithm.forward (and its parent classes) invoke the agent via agent(x), passing exactly one element from the dataset per call .
3. No direct callback function support
   You cannot hand in a raw Python function or lambda to the trainer; it must be wrapped as a trace.model class (or function decorated with @trace.model) to satisfy the Module interface.

Questions

1. Function optimization
   Must every function I want to train be bundled into a trace.model class (or decorated) before passing it to the trainer?

2. Custom workflow steps
   If I need to bypass or extend the agent’s __call__, is the intended pattern to subclass the algorithm (e.g. MinibatchAlgorithm) and override its forward() method to call specific agent methods?

3. Multiple arguments to the agent
   For agents that naturally take more than one input argument, is it better to:

   • Pack them into a single container (e.g. dict or tuple) and pass that as x so I am limited to 1 parameter functions,
   • Or subclass the algorithm and override forward() so it unpacks and passes multiple parameters to the agent directly?

[doxav]

Then, should we need something like this ?

from trace.modules import Module

class CallableModule(Module):
    def __init__(self, fn, unpack='single'):
        self.fn = fn
        self.unpack = unpack

    def forward(self, x):
        if self.unpack == 'tuple':
            return self.fn(*x)
        elif self.unpack == 'dict':
            return self.fn(**x)
        return self.fn(x)


def trace_function(unpack='single'):
    def decorator(fn):
        return CallableModule(fn, unpack=unpack)
    return decorator


# Example usage
@trace_function(unpack='dict')
def my_agent(a, b):
    return a * b

# Used in a trainer (assuming MinibatchAlgorithm is already imported and dataset is defined)
# from trainer.algorithms import MinibatchAlgorithm
# trainer = MinibatchAlgorithm(agent=my_agent, dataset=[{'a': 2, 'b': 4}, {'a': 3, 'b': 5}])
# trainer.train()

[chinganc]

You raised interesting points.

Function optimization
Must every function I want to train be bundled into a trace.model class (or decorated) before passing it to the trainer?

Yes, that's the current design of trainer. We probably can design a basic subclass of Module for training a single callable. Similar to the one you wrote above.

Custom workflow steps
If I need to bypass or extend the agent’s call, is the intended pattern to subclass the algorithm (e.g. MinibatchAlgorithm) and override its forward() method to call specific agent methods?

Yes, the forward is the function to override.

Multiple arguments to the agent
For agents that naturally take more than one input argument, is it better to:

We're thinking about the default as overriding forward which would then contain the right way of using x (a tuple e.g.).

[chinganc]

Since it's common to just optimize a single parameter node or a single function, we can perhaps add a type checking in trainer's __init__ and do some conversion there, so that the user don't have to wrap them into Module.