evaluate_unity.py

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)