Custom Extra Observations in MetaSim (get_extra)#

Design Philosophy#

By default, MetaSim (RoboVerse) provides standard observations such as robot joint states, velocities, and camera images. However, there are situations when you might need specific data that is not directly included in RoboVerse’s default state.

Extra Observation system is a way to inject user-defined observation queries into handlers without modifying the handler class itself.

Architecture#

A custom Extra Observation type is defined by inheriting metasim.queries.base.BaseQueryType . Let’s call the inherited class extra_qeury.

This class has only one important method you need to overwrite call() .

When an Extra Observation is used by a Task, it will be automatically bound to the underlying handler. You can then use self.handler to access the handler instance within the Extra Observation class.

Everytime the handler.get_state() is called, the handler will call the extra_qeury.__call__() method of the Extra Observation type, to gather the extra observations, and return them in the return value of handler.get_state() . When calling handler.get_extra(), the handler will also be explicitly called adn the extra observations will be gathered and returned.

Possible Extra Observations#

Extra Observations includes:

  • Positions or velocities of specific bodies or objects.

  • Sensor readings like IMU data.

  • World-frame poses of specific simulation geometry (e.g., sites, bodies, frames).

  • Any other types you encounter…


Workflow for Using Extra Observations in Tasks#

Adding custom extra observations follows these steps:

  1. Declare the extra observations you need inside your TaskEnv.extra_spec.

  2. Pass these extras to your simulation Handler.

  3. The Handler automatically binds these extras, linking them to your simulation instance.

  4. You then implement the procedure to access the custom data at runtime inside extra_qeury.__call__(). handler.get_states() will use extra_qeury.__call__() to attain extras infomation in states.

Simplified workflow diagram:

Task declares extras ──> Handler binds extras ──> Retrieve extras at runtime

Step-by-Step Usage Guide#

Step 1: Declare Extras in Your Task#

Implement the _extra_spec method inside your task definition:

def _extra_spec(self) -> dict:
    return {
        "imu_pos": SitePos("imu"),   # World position of IMU sensor site
        "head_vel": BodyVel("head"), # Velocity of the robot's head
    }
  • The keys (e.g., "imu_pos") are user-defined and identify the returned data.

  • The values (SitePos, BodyVel) are pre-defined QueryTypes that fetch specific data from your simulator.

Common built-in queries:

  • SitePos("site_name"): Returns world-frame position.

  • BodyVel("body_name"): Returns velocity of a body.


Step 2: Pass Extras to the Handler#

When creating your simulation Handler, pass your extras dictionary (_extra_spec) via the optional_queries parameter:

task = CrawlEnv(CrawlEnv.scenario)
extra_queries = task._extra_spec()

handler = YourSimHandler(
    scenario=task.scenario,
    optional_queries=extra_queries
)

handler.launch()  # Extras get bound here

Step 3: Access Extra Observations at Runtime#

While your simulation is running, simply call:

extras = handler.get_extra()

This returns a dictionary containing your custom observations:

{
    "imu_pos": tensor([[0.1, 0.2, 0.3], [...]]),
    "head_vel": tensor([[0.01, 0.02, 0.03], [...]])
}

You can directly use this data for:

  • Reward computations

  • Logging and analysis

  • Visualization purposes

Additionally, when you call:

state = handler.get_state()

the returned state dictionary automatically includes all the extra observations you’ve defined. This way, your extras seamlessly integrate into RoboVerse’s default simulation state.


Adding Your Own Custom Query#

If the existing queries in MetaSim do not fit your needs, you can define a custom query type by subclassing BaseQueryType.

Each custom query involves two steps:

① Define a subclass of BaseQueryType#

First, create a new Python class that inherits from BaseQueryType. You will typically implement two methods in this class:

  • bind_handler(self, handler): This method is called once when the simulation launches. Here you save necessary references (such as simulator handles, indexes, IDs, or configurations) from your specific simulation handler.

  • __call__(self): This method is invoked at runtime every time you request extras. In this method, you use the previously stored references to retrieve the real-time data you need.

Example:

from metasim.queries.base import BaseQueryType
import torch

class BodyMassQuery(BaseQueryType):
    def bind_handler(self, handler):
        # Store the handler for later use
        self.handler = handler

        # Example: simulator-specific logic
        mod = handler.__class__.__module__
        if mod.startswith("metasim.sim.mujoco"):
            self.mass = handler.physics.model.body_mass  # Example MuJoCo API call
        else:
            raise ValueError(f"Unsupported handler: {mod}")

    def __call__(self):
        # Return the stored body mass as a tensor
        return torch.tensor(self.mass)

② Use your custom query inside _extra_spec#

After defining your custom query class, simply instantiate and use it directly in your task’s _extra_spec method:

def _extra_spec(self) -> dict:
    return {
        "body_mass": BodyMassQuery(), # Your custom query
        "imu_pos": SitePos("imu"),
    }

After following these steps, your handler will automatically bind your custom query, and you can seamlessly retrieve this data during runtime.


Practical Tips#

  • Always ensure simulator-specific code in your custom query is correctly placed inside the bind_handler method. This ensures efficient performance, as binding happens once at simulation launch rather than repeatedly at runtime.

  • Your __call__ method should be fast and efficient, simply retrieving pre-stored references or quickly fetching real-time data.

  • It’s good practice to return data consistently as PyTorch tensors (or similarly structured arrays), as these integrate well with downstream RL pipelines.


Now you’re fully equipped to flexibly extend your RoboVerse simulation with any custom extra observations you need!