CarperAI / CarperAI/trlx

Support logging for non-scalar metrics

Open
#493 2 comments 0 reactions 0 assignees View on GitHub
feature request
Dominant language
Python
Stars
4.8k
Forks
487
PR merge metrics
No merged PRs in 30d

Description

### 🚀 The feature, motivation, and pitch

AccelerateRLTrainer.evaluate() logs a table of generated eval outputs and metrics to the metrics tracker.

If I understand correctly, only scalar metrics are currently supported.

This feature would allow non-scalar metrics to be logged.

#### Use cases:
Allow passthrough logging of non-scalar prompt metadata:

```
def eval_metrics(samples, prompts, outputs, **kwargs):
return kwargs

trlx.train(
model_name,
config=config,
samples=train_samples, # type: ignore
rewards=train_rewards if training_method == "ilql" else None, # type: ignore
eval_prompts=eval_eval_prompts, # type: ignore
metric_fn=eval_metrics,
).model
```

#### Implementation suggestion:

Modify the mean_metrics calculation (below) to only calculate means for values that can be successfully cast to float tensors.

https://github.com/CarperAI/trlx/blob/0dce99d96b7d70b6a9114129d8e38bf6c80eb653/trlx/trainer/accelerate_base_trainer.py#L428-L430

### Alternatives

_No response_

### Additional context

_No response_

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.