def before_hook(step_idx: int, transition: EnvTransition):
    """Called before step processes the transition."""
    print(f"About to execute step {step_idx}")
    # Useful for: logging, validation, setup

def after_hook(step_idx: int, transition: EnvTransition):
    """Called after step has processed the transition."""
    print(f"Completed step {step_idx}")
    # Useful for: monitoring results, cleanup, debugging

processor.register_before_step_hook(before_hook)
processor.register_after_step_hook(after_hook)

def check_nans(step_idx: int, transition: EnvTransition):
    """Check for NaN values in observations."""
    obs = transition.get(TransitionKey.OBSERVATION)
    if obs:
        for key, value in obs.items():
            if isinstance(value, torch.Tensor) and torch.isnan(value).any():
                print(f"NaN detected in {key} at step {step_idx}")

# Register the hook to run after each step
processor.register_after_step_hook(check_nans)

# Process your data - the hook will be called automatically
output = processor(input_data)

# Remove the hook when done debugging
processor.unregister_after_step_hook(check_nans)

class DataProcessorPipeline:
    def __init__(self):
        self.steps = [...]
        self.before_step_hooks = []  # List of before hooks
        self.after_step_hooks = []   # List of after hooks

    def _forward(self, transition):
        """Internal method that processes the transition through all steps."""
        for step_idx, processor_step in enumerate(self.steps):
            # 1. Call all BEFORE hooks
            for hook in self.before_step_hooks:
                hook(step_idx, transition)

            # 2. Execute the actual processing step
            transition = processor_step(transition)

            # 3. Call all AFTER hooks
            for hook in self.after_step_hooks:
                hook(step_idx, transition)

        return transition

    def register_before_step_hook(self, hook_fn):
        self.before_step_hooks.append(hook_fn)

    def register_after_step_hook(self, hook_fn):
        self.after_step_hooks.append(hook_fn)

Input → Before Hook → Step 0 → After Hook → Before Hook → Step 1 → After Hook → ... → Output

def timing_before(step_idx, transition):
    print(f"⏱️  Starting step {step_idx}")

def validation_after(step_idx, transition):
    print(f"✅ Completed step {step_idx}")

processor.register_before_step_hook(timing_before)
processor.register_after_step_hook(validation_after)

# This will output:
# ⏱️  Starting step 0
# ✅ Completed step 0
# ⏱️  Starting step 1
# ✅ Completed step 1
# ⏱️  Starting step 2
# ✅ Completed step 2

def log_shapes(step_idx: int, transition: EnvTransition):
    obs = transition.get(TransitionKey.OBSERVATION)
    if obs:
        print(f"Step {step_idx} observation shapes:")
        for key, value in obs.items():
            if isinstance(value, torch.Tensor):
                print(f"  {key}: {value.shape}")

processor.register_after_step_hook(check_nans)      # Executes first
processor.register_after_step_hook(log_shapes)     # Executes second

# Both hooks will be called after each step in registration order
output = processor(input_data)

# This creates a generator that yields intermediate states.
# The first yield is the input itself; yield i (i >= 1) is the output of step i - 1.
for i, intermediate_result in enumerate(processor.step_through(input_data)):
    print("=== Input ===" if i == 0 else f"=== After step {i - 1} ===")

    # Inspect the observation at this stage
    obs = intermediate_result.get(TransitionKey.OBSERVATION)
    if obs:
        for key, value in obs.items():
            if isinstance(value, torch.Tensor):
                print(f"{key}: shape={value.shape}, dtype={value.dtype}")

# Step through the pipeline with debugging
for i, intermediate in enumerate(processor.step_through(data)):
    print("Input" if i == 0 else f"After step {i - 1}: {processor.steps[i - 1].__class__.__name__}")

    # Set a breakpoint to inspect the current state
    breakpoint()  # Debugger will pause here

    # You can now inspect 'intermediate' in the debugger:
    # - Check tensor shapes and values
    # - Verify expected transformations
    # - Look for unexpected changes

# Get the expected output features from your pipeline
initial_features = {
    PipelineFeatureType.OBSERVATION: {
        "observation.state": PolicyFeature(type=FeatureType.STATE, shape=(7,)),
        "observation.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224))
    },
    PipelineFeatureType.ACTION: {
        "action": PolicyFeature(type=FeatureType.ACTION, shape=(4,))
    }
}

# Check what your pipeline will output
output_features = processor.transform_features(initial_features)

print("Input features:")
for feature_type, features in initial_features.items():
    print(f"  {feature_type}:")
    for key, feature in features.items():
        print(f"    {key}: {feature.type.value}, shape={feature.shape}")

print("\nOutput features:")
for feature_type, features in output_features.items():
    print(f"  {feature_type}:")
    for key, feature in features.items():
        print(f"    {key}: {feature.type.value}, shape={feature.shape}")

# Define what features you expect the pipeline to produce
expected_keys = ["observation.state", "observation.image", "action"]

print("Validating feature contract...")
for expected_key in expected_keys:
    found = False
    for feature_type, features in output_features.items():
        if expected_key in features:
            feature = features[expected_key]
            print(f"✅ {expected_key}: {feature.type.value}, shape={feature.shape}")
            found = True
            break

    if not found:
        print(f"❌ Missing expected feature: {expected_key}")