test_decoder_fix.py

#!/usr/bin/env python3
"""Test script to verify ConvDecoder handles non-square images correctly."""

import sys
sys.path.insert(0, '/Research/dreamerv3-metaworld/dreamerv3-torch-main')

import torch
from networks import ConvDecoder

def test_decoder(shape, feat_size=1024):
    """Test that decoder can handle the given shape."""
    print(f"\n=== Testing decoder with shape {shape} ===")

    try:
        # Create decoder
        decoder = ConvDecoder(feat_size=feat_size, shape=shape, depth=32, act='ELU', minres=4)
        print(f"✓ Decoder created successfully")
        print(f"  Initial dims: h={decoder._minres_h}, w={decoder._minres_w}")
        print(f"  Embed size: {decoder._embed_size}")

        # Create dummy input: batch=6, time=5, features=feat_size
        features = torch.randn(6, 5, feat_size)
        print(f"✓ Created test features with shape {list(features.shape)}")

        # Run forward pass
        output = decoder(features)
        print(f"✓ Forward pass successful")
        print(f"  Output shape: {list(output.shape)}")
        print(f"  Expected shape: [6, 5, {shape[1]}, {shape[2]}, {shape[0]}]")

        # Verify output shape
        expected_shape = [6, 5, shape[1], shape[2], shape[0]]
        if list(output.shape) == expected_shape:
            print(f"✓✓✓ SUCCESS: Output shape matches expected shape!")
            return True
        else:
            print(f"✗✗✗ FAILED: Output shape mismatch!")
            return False

    except Exception as e:
        print(f"✗✗✗ ERROR: {e}")
        import traceback
        traceback.print_exc()
        return False

if __name__ == "__main__":
    print("Testing ConvDecoder with various shapes...")

    # Test 1: Square image (original case)
    test_decoder((3, 64, 64))

    # Test 2: Wide image (3 cameras concatenated horizontally)
    test_decoder((3, 64, 192))

    # Test 3: Another aspect ratio
    test_decoder((3, 64, 128))

    print("\n=== All tests complete ===")