File size: 1,545 Bytes
e027202 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 | """Check the synthetic request against the original training preprocessing."""
import json
from pathlib import Path
import socket
import sys
from flax.traverse_util import flatten_dict
import numpy as np
from pap_runtime import input_transforms, load_request
from openpi import transforms
root = Path(__file__).resolve().parent.parent
def no_network(*args, **kwargs):
raise AssertionError('Offline inference preprocessing attempted a network connection')
socket.socket.connect = no_network
obs = load_request(root / 'runtime/examples/synthetic_request.npz')
data = {'state':obs['state'], 'actions':obs['action_prefix'], 'images':obs['images'], 'prompt':obs['prompt']}
out = flatten_dict(transforms.compose(input_transforms(root))(data))
errors = {}
with np.load(root / 'runtime/examples/preprocessed_reference.npz', allow_pickle=False) as reference:
assert set(reference.files) == {'/'.join(k) for k in out}
for key, value in out.items():
name = '/'.join(key)
np.testing.assert_allclose(value, reference[name], rtol=0, atol=1e-7)
errors[name] = float(np.max(np.abs(np.asarray(value,dtype=float)-reference[name].astype(float))))
assert 'openpi.training.config' not in sys.modules
assert 'openpi.training.data_loader' not in sys.modules
import openpi.models.pi0
assert str(Path(openpi.models.pi0.__file__).resolve()).startswith(str(root / 'runtime/src'))
print(json.dumps({'status':'passed','offline':True,'training_config_imported':False,
'preprocessing_max_errors':errors},indent=2))
|