Implementation:Google deepmind Dm control Hooks Test Utils

Knowledge Sources	Google_deepmind_Dm_control
Domains	Robotics Simulation, Testing
Last Updated	2026-02-15 04:00 GMT

Overview

This module provides test utilities for verifying the correct ordering and call counts of Composer lifecycle hooks (callbacks) during environment episodes.

Description

The hooks_test_utils module implements a set of classes that enable rigorous testing of the Composer environment's callback dispatch system. The central class, HooksTracker, is a mixin that implements all Composer lifecycle callbacks -- initialize_episode_mjcf, after_compile, initialize_episode, before_step, before_substep, after_substep, and after_step -- while tracking call counts and asserting correct ordering invariants at each invocation.

For example, HooksTracker verifies that before_step is only called after all initialization hooks have completed, that physics step counts match before_substep counts, and that substep counts align properly with control step counts. The tracker uses MRO inspection to determine whether to call super() methods, enabling it to work as a mixin with both Entity and Task classes.

TrackedEntity combines HooksTracker with composer.Entity, while TrackedTask combines it with composer.NullTask. The HooksTestMixin class sets up a 9-entity hierarchy in a tree structure (root entity 0 with children 1, 2, 3; entity 1 has children 4, 5; entity 2 has children 6, 7; entity 4 has child 8) and provides a track_episode context manager that enables tracking, runs the episode, and asserts all hooks were called the expected number of times.

The helper function add_bodies_and_actuators adds body/joint/actuator elements to an MJCF model for testing purposes.

Usage

Use this module in unit tests that need to verify the correct ordering and frequency of Composer lifecycle hook calls. It is used as the backbone for environment_hooks_test.py and similar test files that validate the callback scheduling contract of the Composer framework.

Code Reference

Source Location

• Repository: Google_deepmind_Dm_control
• File: dm_control/composer/hooks_test_utils.py
• Lines: 1-323

Signature

def add_bodies_and_actuators(mjcf_model, num_actuators): ...

class HooksTracker:
    def __init__(self, test_case, physics_timestep, control_timestep, *args, **kwargs): ...
    def assertEqual(self, actual, expected, msg=''): ...
    def assertHooksNotCalled(self, *hook_names): ...
    def assertHooksCalledOnce(self, *hook_names): ...
    def assertCompleteEpisode(self, control_steps): ...
    def assertPhysicsStepCountEqual(self, physics, expected_count): ...
    def reset_call_counts(self): ...
    # Composer lifecycle callbacks:
    def initialize_episode_mjcf(self, random_state): ...
    def after_compile(self, physics, random_state): ...
    def initialize_episode(self, physics, random_state): ...
    def before_step(self, physics, *args): ...
    def before_substep(self, physics, *args): ...
    def after_substep(self, physics, random_state): ...
    def after_step(self, physics, random_state): ...

class TrackedEntity(HooksTracker, composer.Entity):
    def _build(self, name): ...
    mjcf_model -> property
    name -> property

class TrackedTask(HooksTracker, composer.NullTask):
    def __init__(self, physics_timestep, control_timestep, *args, **kwargs): ...

class HooksTestMixin:
    def setUp(self): ...
    def track_episode(self): ...  # context manager

Import

from dm_control.composer import hooks_test_utils

I/O Contract

Inputs

Name	Type	Required	Description
test_case	absltest.TestCase	Yes	The test case instance that provides assertion methods
physics_timestep	float	Yes	The physics simulation timestep in seconds
control_timestep	float	Yes	The control loop timestep in seconds
num_actuators	int	Yes (for add_bodies_and_actuators)	Number of actuators to add (must be a multiple of 2)

Outputs

Name	Type	Description
tracked	bool	Whether the tracker is currently monitoring hook calls
assertions	None	Raises AssertionError if hook ordering invariants are violated

Usage Examples

Basic Usage

from absl.testing import absltest
from dm_control import composer
from dm_control.composer import hooks_test_utils

class MyHooksTest(hooks_test_utils.HooksTestMixin, absltest.TestCase):

    def test_hooks_are_called_correctly(self):
        env = composer.Environment(self.task)
        with self.track_episode():
            env.reset()
            for _ in range(self.steps_per_episode):
                env.step(env.action_spec().generate_value())