"""Test whether speculative rollout state can affect the next executed step."""
import os,sys,json
from pathlib import Path
ROOT=Path(__file__).resolve().parents[1];O=ROOT/'output/spider_dynamics_fix_20260915';T=O/'datasets/processed/current/l20/bimanual/boxes'
os.environ['WARP_CACHE_PATH']=str(ROOT/'.spider_cache/warp');sys.path.insert(0,str(ROOT/'third_party/spider'))
import numpy as np,torch,mujoco,warp as wp
from spider.config import Config
from spider.simulators import mjwp as s
c=Config(device='cuda:0',num_samples=64,nconmax_per_env=1024,njmax_per_env=3072,sim_dt=.0025,embodiment_type='bimanual');c.model_path=str(T/'scene_act.xml');c.object_actuator_ids=list(range(44,56));c.npair=0
setup=s.setup_mj_model
def model(c):
 m=setup(c);m.opt.iterations=80;return m
s.setup_mj_model=model
r=np.load(T/'0/trajectory_kinematic_act.npz');ref=tuple(torch.tensor(r[k],device='cuda:0',dtype=torch.float32) for k in ['qpos','qvel','ctrl','contact','contact_pos']);env=s.setup_env(c,ref);u=ref[2][0];state=s.save_state(env);s.step_env(c,env,u);wp.synchronize();expected=env.data_wp.qpos.numpy().copy();s.load_state(env,state)
for _ in range(80):s.step_env(c,env,u)
s.load_state(env,state);s.step_env(c,env,u);wp.synchronize();after=env.data_wp.qpos.numpy().copy();report={'state_only_max_qpos_difference':float(np.max(abs(expected-after)))}
s.load_state(env,state);params=s.save_env_params(c,env);s.load_env_params(c,env,{'kp':np.ones(12,dtype=np.float32)*10,'kd':np.ones(12,dtype=np.float32)})
for _ in range(80):s.step_env(c,env,u)
s.load_state(env,state);s.load_env_params(c,env,params);s.step_env(c,env,u);wp.synchronize();after=env.data_wp.qpos.numpy().copy();report['state_and_parameters_max_qpos_difference']=float(np.max(abs(expected-after)));(O/'state_restore_check.json').write_text(json.dumps(report,indent=2));print(report)