add-functor

name	add-functor
description	Use when adding a new observation, event, reward, action, dataset, or randomization functor to an EmbodiChain environment

Add Functor

Scaffold a new functor following EmbodiChain's Functor/FunctorCfg pattern.

When to Use

• User asks to add an observation term, reward function, event handler, action term, dataset functor, or randomizer
• User says "add a reward", "new observation", "create a randomizer", "add event functor"
• Any new function needs to be registered in a manager config

Determine Functor Type

Functor Type	Config Class	Module File	Manager	Signature
Observation	ObservationCfg (extends FunctorCfg)	managers/observations.py	ObservationManager	(env, obs, entity_cfg, ...) -> Tensor
Reward	RewardCfg (extends FunctorCfg)	managers/rewards.py	RewardManager	(env, obs, action, info, ...) -> Tensor
Event	EventCfg (extends FunctorCfg)	managers/events.py	EventManager	(env, env_ids, ...) -> None
Action	ActionTermCfg (extends FunctorCfg)	managers/actions.py	ActionManager	Varies
Dataset	DatasetFunctorCfg (extends FunctorCfg)	managers/datasets.py	DatasetManager	(env, ...) -> dict
Randomization	EventCfg (randomizations ARE events)	managers/randomization/<type>.py	EventManager	(env, env_ids, entity_cfg, ...) -> None

Two Functor Styles

Function-style (Preferred for Simple Functors)

A plain function with the right signature. Registered via FunctorCfg(func=my_function, params={...}).

def my_reward(
    env: EmbodiedEnv,
    obs: dict,
    action: EnvAction,
    info: dict,
    my_param: float = 1.0,       # params become keyword args
) -> torch.Tensor:
    """Short one-line summary.

    Longer description if needed.

    Args:
        env: The environment instance.
        obs: The observation dictionary.
        action: The action taken.
        info: The info dictionary.
        my_param: Description of this parameter.

    Returns:
        Reward tensor of shape (num_envs,).
    """
    # implementation
    return result

Class-style (Required When Functor Has State)

A class inheriting Functor, with __init__(cfg, env) and __call__(env, ...). Registered via FunctorCfg(func=MyClass, params={...}).

class my_randomizer(Functor):
    """One-line summary."""

    def __init__(self, cfg: FunctorCfg, env: EmbodiedEnv):
        super().__init__(cfg, env)
        # Extract params and initialize state
        self.entity_cfg: SceneEntityCfg = cfg.params["entity_cfg"]

    def __call__(self, env: EmbodiedEnv, env_ids: torch.Tensor, **kwargs):
        """Apply the randomization.

        Args:
            env: The environment instance.
            env_ids: Target environment IDs.
        """
        # implementation

Steps

1. Identify Functor Type and Style

Ask the user:

1. Which manager? (observation / reward / event / action / dataset / randomization)
2. Function or class style? (function for stateless, class for stateful)
3. What does it do? (brief description for naming + docstring)

2. Choose the Right Module File

Place the functor in the existing module for its type:

Type	File
Observation	embodichain/lab/gym/envs/managers/observations.py
Reward	embodichain/lab/gym/envs/managers/rewards.py
Event	embodichain/lab/gym/envs/managers/events.py
Action	embodichain/lab/gym/envs/managers/actions.py
Dataset	embodichain/lab/gym/envs/managers/datasets.py
Physics randomization	embodichain/lab/gym/envs/managers/randomization/physics.py
Visual randomization	embodichain/lab/gym/envs/managers/randomization/visual.py
Spatial randomization	embodichain/lab/gym/envs/managers/randomization/spatial.py
Geometry randomization	embodichain/lab/gym/envs/managers/randomization/geometry.py

3. Write the Functor

Follow the template for function-style or class-style (see above).

Key rules:

• First argument is always env: EmbodiedEnv (use TYPE_CHECKING guard for the import)
• Use from __future__ import annotations at the top
• Use SceneEntityCfg for entity references, not raw strings
• For observation functors: add shape key to FunctorCfg.extra dict
• For randomization functors: second arg is env_ids: torch.Tensor | list[int]
• For reward functors: return shape must be (num_envs,)

4. Update __all__

Add the new functor to the module's __all__ list. If no __all__ exists, create one.

5. Write a Test

Place at tests/gym/envs/managers/test_<functor_type>.py (append to existing file if present).

For functors that don't need a live simulation, use mock objects (MockEnv, MockSim, etc.) following the pattern in tests/gym/envs/managers/test_reward_functors.py.

6. Run black

black embodichain/lab/gym/envs/managers/<module>.py
black tests/gym/envs/managers/test_<functor_type>.py

Common Mistakes

Mistake	Fix
Wrong first argument signature	Observation: (env, obs, ...), Reward: (env, obs, action, info, ...), Event/Randomization: (env, env_ids, ...)
Importing EmbodiedEnv at module level	Use TYPE_CHECKING guard to avoid circular imports
Forgetting SceneEntityCfg for entity refs	Always use SceneEntityCfg(uid="...") not bare strings
Returning wrong tensor shape	Rewards must return (num_envs,), observations must match declared shape
Missing from __future__ import annotations	Required in every file
Class-style functor not calling super().__init__	Always call super().__init__(cfg, env)
Adding randomizer as standalone	Randomizations ARE events — they go in randomization/ but use EventCfg

Quick Reference

Step	Action
1	Identify manager type + function vs class style
2	Write functor in the correct module file
3	Update __all__ in that module
4	Write test with mocks (no sim needed for most)
5	Run black on changed files