| import pathlib,json,numpy as np,onnxruntime as ort |
| from mlagents_envs.environment import UnityEnvironment |
| from mlagents_envs.base_env import ActionTuple |
| from mlagents_envs.side_channel.engine_configuration_channel import EngineConfigurationChannel |
| root=pathlib.Path('/content/hf-rl') |
| for number,name in enumerate(['SnowballTarget','Pyramids']): |
| binary=next(p for p in (root/name).rglob('*') if p.is_file() and p.name in [name,name+'.x86_64'] and p.read_bytes()[:4]==bytes([127,69,76,70])) |
| channel=EngineConfigurationChannel();channel.set_configuration_parameters(time_scale=20) |
| env=UnityEnvironment(file_name=str(binary),seed=100001+number,worker_id=10+number,no_graphics=True,side_channels=[channel]) |
| env.reset();behavior=next(iter(env.behavior_specs));spec=env.behavior_specs[behavior] |
| session=ort.InferenceSession(str(root/'results'/('Kay-'+name)/(name+'.onnx'))) |
| totals={};rewards=[] |
| for step in range(100000): |
| decisions,terminals=env.get_steps(behavior) |
| for i,aid in enumerate(terminals.agent_id): |
| rewards.append(float(totals.pop(int(aid),0)+terminals.reward[i])) |
| if len(rewards)>=100:break |
| if len(decisions): |
| for i,aid in enumerate(decisions.agent_id):totals[int(aid)]=totals.get(int(aid),0)+float(decisions.reward[i]) |
| feed={f'obs_{i}':x.astype(np.float32) for i,x in enumerate(decisions.obs)} |
| feed['action_masks']=np.ones((len(decisions),sum(spec.action_spec.discrete_branches)),dtype=np.float32) if decisions.action_mask is None else 1-np.concatenate(decisions.action_mask,axis=1).astype(np.float32) |
| actions=session.run(['deterministic_discrete_actions'],feed)[0].astype(np.int32) |
| env.set_actions(behavior,ActionTuple(discrete=actions)) |
| env.step() |
| env.close() |
| report={'environment':'ML-Agents-'+name,'evaluation_seed':100001+number,'episodes':len(rewards),'mean_reward':float(np.mean(rewards)),'std_reward':float(np.std(rewards)),'episode_returns':rewards,'policy':'deterministic ONNX actions','training_seed':42} |
| (root/'results'/('Kay-'+name)/'evaluation.json').write_text(json.dumps(report,indent=2)) |
| print({k:v for k,v in report.items() if k!='episode_returns'},flush=True) |
|
|
