eval_policy.py

# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

importargparse
importwarnings

importnumpyasnp

fromgr00t.data.datasetimportLeRobotSingleDataset
fromgr00t.eval.robotimportRobotInferenceClient
fromgr00t.experiment.data_configimportDATA_CONFIG_MAP
fromgr00t.model.policyimportBasePolicy, Gr00tPolicy
fromgr00t.utils.evalimportcalc_mse_for_single_trajectory

warnings.simplefilter("ignore", category=FutureWarning)

"""
Example command:

python scripts/eval_policy.py --host localhost --port 5555 --plot
 --modality_keys right_arm right_hand
 --steps 250
 --trajs 1000
 --action_horizon 16
 --video_backend decord
 --dataset_path demo_data/robot_sim.PickNPlace/
 --embodiment_tag gr1
 --data_config gr1_arms_waist
provide --model_path to load up the model checkpoint in this script.
"""

if__name__=="__main__":
parser=argparse.ArgumentParser()
parser.add_argument("--host", type=str, default="localhost", help="host")
parser.add_argument("--port", type=int, default=5555, help="port")
parser.add_argument("--plot", action="store_true", help="plot images")
parser.add_argument("--modality_keys", nargs="+", type=str, default=["right_arm", "right_hand"])
parser.add_argument(
"--data_config",
type=str,
default="gr1_arms_only",
choices=list(DATA_CONFIG_MAP.keys()),
help="data config name",
 )
parser.add_argument("--steps", type=int, default=150, help="number of steps to run")
parser.add_argument("--trajs", type=int, default=1, help="trajectories to run")
parser.add_argument("--action_horizon", type=int, default=16)
parser.add_argument("--video_backend", type=str, default="decord")
parser.add_argument("--dataset_path", type=str, default="demo_data/robot_sim.PickNPlace/")
parser.add_argument(
"--embodiment_tag",
type=str,
help="The embodiment tag for the model.",
default="gr1",
 )
## When using a model instead of client-server mode.
parser.add_argument(
"--model_path",
type=str,
default=None,
help="[Optional] Path to the model checkpoint directory, this will disable client server mode.",
 )
parser.add_argument(
"--denoising_steps",
type=int,
help="Number of denoising steps if model_path is provided",
default=4,
 )
args=parser.parse_args()

data_config=DATA_CONFIG_MAP[args.data_config]
ifargs.model_pathisnotNone:
importtorch

modality_config=data_config.modality_config()
modality_transform=data_config.transform()

policy: BasePolicy=Gr00tPolicy(
model_path=args.model_path,
modality_config=modality_config,
modality_transform=modality_transform,
embodiment_tag=args.embodiment_tag,
denoising_steps=args.denoising_steps,
device="cuda"iftorch.cuda.is_available() else"cpu",
 )
else:
policy: BasePolicy=RobotInferenceClient(host=args.host, port=args.port)

all_gt_actions= []
all_pred_actions= []

# Get the supported modalities for the policy
modality=policy.get_modality_config()
print(modality)

# Create the dataset
dataset=LeRobotSingleDataset(
dataset_path=args.dataset_path,
modality_configs=modality,
video_backend=args.video_backend,
video_backend_kwargs=None,
transforms=None, # We'll handle transforms separately through the policy
embodiment_tag=args.embodiment_tag,
 )

print(len(dataset))
# Make a prediction
obs=dataset[0]
fork, vinobs.items():
ifisinstance(v, np.ndarray):
print(k, v.shape)
else:
print(k, v)

fork, vindataset.get_step_data(0, 0).items():
ifisinstance(v, np.ndarray):
print(k, v.shape)
else:
print(k, v)

print("Total trajectories:", len(dataset.trajectory_lengths))
print("All trajectories:", dataset.trajectory_lengths)
print("Running on all trajs with modality keys:", args.modality_keys)

all_mse= []
fortraj_idinrange(args.trajs):
print("Running trajectory:", traj_id)
mse=calc_mse_for_single_trajectory(
policy,
dataset,
traj_id,
modality_keys=args.modality_keys,
steps=args.steps,
action_horizon=args.action_horizon,
plot=args.plot,
 )
print("MSE:", mse)
all_mse.append(mse)
print("Average MSE across all trajs:", np.mean(all_mse))
print("Done")
exit()